From c21d55bd1b202d47df5713532642323ccb946fcb Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 16 Apr 2026 13:49:48 +0200 Subject: [PATCH 01/89] Benchmark the mcq medeical dataset --- cli/package.json | 2 + cli/src/evaluate_finetuned_gpt2.ts | 243 ++++++++++++++++++++++++ discojs/src/default_tasks/index.ts | 3 +- discojs/src/default_tasks/privacyrun.ts | 66 +++++++ 4 files changed, 313 insertions(+), 1 deletion(-) create mode 100644 cli/src/evaluate_finetuned_gpt2.ts create mode 100644 discojs/src/default_tasks/privacyrun.ts diff --git a/cli/package.json b/cli/package.json index cc0f741e2..8ed779b3f 100644 --- a/cli/package.json +++ b/cli/package.json @@ -9,6 +9,8 @@ "benchmark_gpt": "npm run build && node dist/benchmark_gpt.js", "train_gpt": "npm run build && node dist/train_gpt.js", "hellaswag_gpt": "npm run build && node dist/hellaswag_gpt.js", + "eval_finetuned_gpt2": "npm run build && node dist/evaluate_finetuned_gpt2.js", + "finetune_gpt": "npm run build && node dist/finetune_gpt.js", "build": "tsc --build", "test": ": nothing" }, diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts new file mode 100644 index 000000000..6f031cd04 --- /dev/null +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -0,0 +1,243 @@ +import "@tensorflow/tfjs-node"; +import * as tf from "@tensorflow/tfjs"; +import fs from "node:fs/promises"; +import { parse } from "ts-command-line-args"; +import { models, Tokenizer } from "@epfml/discojs"; +import { loadModelFromDisk } from "@epfml/discojs-node"; + +interface Args { + modelPath: string; + testPath: string; + maxSamples?: number; + savePath?: string; + help?: boolean; +} + +// ========================= +// HOW TO RUN +// ========================= +// npm -w cli run eval_finetuned_gpt2 -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/train_no_exp.txt --maxSamples 100 + +// ========================= +// LOAD DATASET +// ========================= +async function loadDataset(filePath: string, limit = -1): Promise { + const text = await fs.readFile(filePath, "utf-8"); + const lines = text.split("\n"); + + const samples: string[] = []; + let current = ""; + + for (const line of lines) { + const l = line.trim(); + + if (l.includes("<|startoftext|>")) { + current = ""; + } else if (l.includes("<|endoftext|>")) { + samples.push(current.trim()); + if (limit !== -1 && samples.length >= limit) break; + } else { + current += l + "\n"; + } + } + + return samples; +} + +// ========================= +// PARSE SAMPLE +// ========================= +function parseSample(sample: string) { + const lines = sample.split("\n"); + + let answer = ""; + const promptLines: string[] = []; + + for (const line of lines) { + if (line.startsWith("Answer:")) { + answer = line.replace("Answer:", "").trim(); + } else { + promptLines.push(line); + } + } + + const basePrompt = promptLines.join("\n"); + return { basePrompt, answer }; +} + +// ========================= +// SOFTMAX (for safety) +// ========================= +async function scoreText( + tfModel: tf.LayersModel, + tokenizer: Tokenizer, + text: string +): Promise { + const tokens = tokenizer.tokenize(text); + + if (tokens.size < 2) return -Infinity; + + const inputTokens = tokens.slice(0, tokens.size - 1).toArray(); + const targets = tokens.slice(1).toArray(); + + const inputTensor = tf.tensor([inputTokens], [1, inputTokens.length], "int32"); + + const logits = tfModel.predict(inputTensor) as tf.Tensor; + const logitsArray = await logits.array() as number[][][]; + + let score = 0; + + for (let i = 0; i < targets.length; i++) { + const stepLogits = logitsArray[0][i]; + + const logit = stepLogits[targets[i]] ?? -100; + + score += logit; + } + + inputTensor.dispose(); + logits.dispose(); + + return score; +} + +// ========================= +// SCORE OPTIONS +// ========================= +async function scoreOptions( + tfModel: tf.LayersModel, + tokenizer: Tokenizer, + texts: string[] +): Promise { + const scores: number[] = []; + + for (const t of texts) { + const s = await scoreText(tfModel, tokenizer, t); + scores.push(s); + } + + return scores; +} + +// ========================= +// BENCHMARK +// ========================= +async function benchmarkQA( + model: models.GPT, + tokenizer: Tokenizer, + dataset: string[], + savePath?: string +) { + console.log("=== QA LOGPROB BENCHMARK ==="); + + const tfModel = model.extract(); + + let correct = 0; + let total = 0; + + const options = ["A", "B", "C", "D"]; + + const confusion: Record> = { + A: { A: 0, B: 0, C: 0, D: 0 }, + B: { A: 0, B: 0, C: 0, D: 0 }, + C: { A: 0, B: 0, C: 0, D: 0 }, + D: { A: 0, B: 0, C: 0, D: 0 } + }; + + const logs: any[] = []; + + const start = Date.now(); + + for (const sample of dataset) { + const { basePrompt, answer } = parseSample(sample); + + const texts = options.map( + (opt) => `${basePrompt}\nAnswer: ${opt}` + ); + + const scores = await scoreOptions(tfModel, tokenizer, texts); + + let bestIdx = 0; + for (let i = 1; i < scores.length; i++) { + if (scores[i] > scores[bestIdx]) bestIdx = i; + } + + const predicted = options[bestIdx]; + + if (predicted === answer) correct++; + total++; + + if (confusion[answer]) { + confusion[answer][predicted]++; + } + + logs.push({ + predicted, + answer, + correct: predicted === answer + }); + + if (total % 50 === 0) { + console.log(`Processed ${total} samples...`); + } + } + + const accuracy = correct / total; + const duration = ((Date.now() - start) / 1000).toFixed(2); + + console.log("\n========================="); + console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); + console.log(`Time: ${duration}s`); + console.log("=========================\n"); + + console.log("Confusion Matrix:"); + console.table(confusion); + + console.log("\nPer-class accuracy:"); + for (const cls of options) { + const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); + const correctCls = confusion[cls][cls]; + const acc = totalCls ? (correctCls / totalCls) * 100 : 0; + + console.log(`${cls}: ${acc.toFixed(2)}%`); + } + + if (savePath) { + await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); + console.log(`Saved results to ${savePath}`); + } +} + +// ========================= +// MAIN +// ========================= +async function main() { + const args = parse({ + modelPath: { type: String }, + testPath: { type: String }, + maxSamples: { type: Number, optional: true, defaultValue: 100 }, + savePath: { type: String, optional: true }, + help: { type: Boolean, optional: true } + }); + + console.log("Loading tokenizer..."); + const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + + console.log("Loading model..."); + const model = await loadModelFromDisk(args.modelPath); + + if (!(model instanceof models.GPT)) { + throw new Error("Model must be GPT"); + } + + console.log("Loading dataset..."); + const dataset = await loadDataset(args.testPath, args.maxSamples); + + console.log(`Loaded ${dataset.length} samples`); + + await benchmarkQA(model, tokenizer, dataset, args.savePath); + + console.log("Done."); +} + +main().catch(console.error); \ No newline at end of file diff --git a/discojs/src/default_tasks/index.ts b/discojs/src/default_tasks/index.ts index 43adf0d3c..f46f27026 100644 --- a/discojs/src/default_tasks/index.ts +++ b/discojs/src/default_tasks/index.ts @@ -4,4 +4,5 @@ export { mnist } from './mnist.js' export { simpleFace } from './simple_face.js' export { titanic } from './titanic.js' export { wikitext } from './wikitext.js' -export { tinderDog } from './tinder_dog.js' \ No newline at end of file +export { tinderDog } from './tinder_dog.js' +export { privacyrun } from './privacyrun.js' \ No newline at end of file diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts new file mode 100644 index 000000000..98dc52aad --- /dev/null +++ b/discojs/src/default_tasks/privacyrun.ts @@ -0,0 +1,66 @@ +import type { TaskProvider } from "../index.js"; +import { Tokenizer, models, serialization } from "../index.js"; + +export const privacyrun: TaskProvider<"text", "local"> = { + async getTask() { + return { + id: 'privacyrun_task', + dataType: "text", + displayInformation: { + title: "GPT Privacy-Preserving Fine-tuning", + summary: { + preview: 'Fine-tune a pre-trained GPT model collaboratively and privately.', + overview: "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning." + }, + model: [ + "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", + "The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer.", + "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", + "Context length is kept at 1024 to match the pre-trained model, with batch size at 1.", + ].join(" "), + dataFormatInformation: 'You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.', + dataExample: + "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", + }, + trainingInformation: { + scheme: 'local', + aggregationStrategy: 'mean', + minNbOfParticipants: 2, + epochs: 6, + validationSplit: 0.1, + roundDuration: 2, + batchSize: 8, + tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), + contextLength: 1024, + tensorBackend: 'gpt' + } + } + }, + + async getModel() { + // Load the pre-trained ONNX-converted model from Google Cloud Storage + // The model should be in DiscoJS serialization format (created by onnx-converter) + const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model.json"; + + try { + const response = await fetch(modelUrl); + if (!response.ok) { + throw new Error(`HTTP ${response.status}: ${response.statusText}`); + } + + const arrayBuffer = await response.arrayBuffer(); + const encodedData = new Uint8Array(arrayBuffer); + + const model = await serialization.model.decode(encodedData); + + if (!(model instanceof models.GPT)) { + throw new Error("Loaded model is not a GPT model"); + } + + return model; + } catch (error) { + console.error("Failed to load model from Google Cloud Storage:", error); + throw new Error(`Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`); + } + }, +} From 3014cf53fe128b8d5490406b584736fa39b0ce66 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 17 Apr 2026 19:50:21 +0200 Subject: [PATCH 02/89] lint error correction --- cli/src/evaluate_finetuned_gpt2.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 6f031cd04..ecccc6bc1 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -144,7 +144,13 @@ async function benchmarkQA( D: { A: 0, B: 0, C: 0, D: 0 } }; - const logs: any[] = []; + type PredictionLog = { + predicted: string; + answer: string; + correct: boolean; + }; + + const logs: PredictionLog[] = []; const start = Date.now(); From e785cf64162209fd7f4e52c8c1488da34b8c8d6a Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Apr 2026 13:57:13 +0200 Subject: [PATCH 03/89] add val dataset path param --- cli/src/args.ts | 11 +++++-- cli/src/cli.ts | 15 +++++++-- cli/src/data.ts | 10 ++++-- discojs/src/default_tasks/privacyrun.ts | 6 ++-- discojs/src/training/disco.ts | 44 +++++++++++++++++++++---- 5 files changed, 70 insertions(+), 16 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index ced893a72..955dcf47e 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -21,6 +21,8 @@ export interface BenchmarkArguments { roundDuration: number batchSize: number validationSplit: number + datasetPath?: string + validationDatasetPath?: string // DP epsilon?: number @@ -41,6 +43,8 @@ export interface BenchmarkArguments { type BenchmarkUnsafeArguments = Omit & { task: string + datasetPath?: string + validationDatasetPath?: string help?: boolean } @@ -55,6 +59,8 @@ const unsafeArgs = parse( roundDuration: { type: Number, alias: 'r', description: 'Round duration (in epochs)', defaultValue: 2 }, batchSize: { type: Number, alias: 'b', description: 'Training batch size', defaultValue: 10 }, validationSplit : { type: Number, alias: 'v', description: 'Validation dataset ratio', defaultValue: 0.2 }, + 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 }, save: { type: Boolean, alias: 's', description: 'Save logs of benchmark', defaultValue: false }, host: { type: (raw: string) => new URL(raw), @@ -89,18 +95,19 @@ const unsafeArgs = parse( const supportedTasks = Map( await Promise.all( - Set.of>( + Set.of>( defaultTasks.cifar10, defaultTasks.lusCovid, defaultTasks.simpleFace, defaultTasks.titanic, defaultTasks.tinderDog, defaultTasks.mnist, + defaultTasks.privacyrun, ).map( async (t) => [(await t.getTask()).id, t] as [ string, - TaskProvider<"image" | "tabular", Network>, + TaskProvider<"image" | "tabular" | "text", Network>, ], ), ), diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 2e23c6514..9629f5b01 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -17,6 +17,7 @@ import type { } from "@epfml/discojs"; import { Disco, aggregator as aggregators, client as clients } from '@epfml/discojs' +import { loadText } from "@epfml/discojs-node"; import { getTaskData } from './data.js' import { args } from './args.js' import { makeUserLogFile } from "./user_log.js"; @@ -27,6 +28,7 @@ async function runUser( task: Task, url: URL, data: Dataset, + validationData: Dataset | undefined, userIndex: number, numberOfUsers: number, ): Promise> { @@ -49,7 +51,7 @@ async function runUser( } try{ - for await (const log of disco.trainSummary(data)){ + for await (const log of disco.trainSummary(data, validationData)){ finalLog.push(log); if (jsonStream){ @@ -104,10 +106,17 @@ async function main( console.log({ args }) const dataSplits = await Promise.all( - Range(0, numberOfUsers).map(async i => getTaskData(task.id, i, numberOfUsers)) + Range(0, numberOfUsers).map(async i => getTaskData(task.id, i, numberOfUsers, args.datasetPath)) ) + + let validationData: Dataset | undefined = undefined; + if (args.validationDatasetPath) { + // Assume text task for now + validationData = loadText(args.validationDatasetPath).cached() as Dataset; + } + const logs = await Promise.all( - dataSplits.map((data, i) => runUser(task, args.host, data as Dataset, i, numberOfUsers)) + dataSplits.map((data, i) => runUser(task, args.host, data as Dataset, validationData, i, numberOfUsers)) ) if (args.save) { diff --git a/cli/src/data.ts b/cli/src/data.ts index aa4d0a330..f3b6ff834 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -5,8 +5,9 @@ import { DataType, Image, Task, + Text, } from "@epfml/discojs"; -import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node"; +import { loadCSV, loadImage, loadImagesInDir, loadText } from "@epfml/discojs-node"; import { Repeat } from "immutable"; async function loadSimpleFaceData(userIdx: number, totalClient: number): Promise> { @@ -94,7 +95,10 @@ function loadData(dataName: string, split: number): Dataset( taskID: Task.ID, userIdx: number, - totalClient: number + totalClient: number, + datasetPath?: string, + isValidation?: boolean, + validationDatasetPath?: string ): Promise> { switch (taskID) { case "simple_face": // remove @@ -118,6 +122,8 @@ export async function getTaskData( case "mnist_federated": case "mnist": return loadData("mnist", userIdx) as Dataset; + case "privacyrun": + return loadText(isValidation && validationDatasetPath ? validationDatasetPath : datasetPath ?? '../datasets/med_mcq/train.txt') as Dataset; default: throw new Error(`Data loader for ${taskID} not implemented.`); } diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index 98dc52aad..9b54af387 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -4,7 +4,7 @@ import { Tokenizer, models, serialization } from "../index.js"; export const privacyrun: TaskProvider<"text", "local"> = { async getTask() { return { - id: 'privacyrun_task', + id: 'privacyrun', dataType: "text", displayInformation: { title: "GPT Privacy-Preserving Fine-tuning", @@ -25,7 +25,7 @@ export const privacyrun: TaskProvider<"text", "local"> = { trainingInformation: { scheme: 'local', aggregationStrategy: 'mean', - minNbOfParticipants: 2, + // minNbOfParticipants: 2, epochs: 6, validationSplit: 0.1, roundDuration: 2, @@ -56,6 +56,8 @@ export const privacyrun: TaskProvider<"text", "local"> = { if (!(model instanceof models.GPT)) { throw new Error("Loaded model is not a GPT model"); } + + console.log("Successfully loaded pre-trained GPT model from Google Cloud Storage"); return model; } catch (error) { diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 7dd019b2d..7d85ba8fb 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -159,16 +159,18 @@ export class Disco extends EventEmitter<{ /** Train on dataset, yielding logs of every batch. */ async *trainByBatch( dataset: Dataset, + validationDataset?: Dataset, ): AsyncGenerator { - for await (const round of this.train(dataset)) + for await (const round of this.train(dataset, validationDataset)) for await (const epoch of round) yield* epoch; } /** Train on dataset, yielding summary logs */ async *trainSummary( dataset: Dataset, + validationDataset?: Dataset, ): AsyncGenerator { - for await (const [roundNum, round] of enumerate(this.train(dataset))) { + for await (const [roundNum, round] of enumerate(this.train(dataset, validationDataset))) { const [roundGen, roundLogsPromise] = async_iterator.split(round); const epochResults: Array<{epochNum: number; epochLogs: EpochLogs}> = []; @@ -190,8 +192,8 @@ export class Disco extends EventEmitter<{ } /** Run whole train on dataset. */ - async trainFully(dataset: Dataset): Promise { - for await (const round of this.train(dataset)) + async trainFully(dataset: Dataset, validationDataset?: Dataset): Promise { + for await (const round of this.train(dataset, validationDataset)) for await (const epoch of round) for await (const _ of epoch); } @@ -203,20 +205,23 @@ export class Disco extends EventEmitter<{ **/ async *train( dataset: Dataset, + validationDataset?: Dataset, ): AsyncGenerator< AsyncGenerator, RoundLogs> > { this.#logger.success("Training started"); - const [trainingDataset, validationDataset] = - await this.#preprocessSplitAndBatch(dataset); + const [trainingDataset, validationDataset_] = + validationDataset !== undefined + ? await this.#preprocessDatasets(dataset, validationDataset) + : await this.#preprocessSplitAndBatch(dataset); // the client fetches the latest weights upon connection // TODO unsafe cast this.trainer.model = (await this.#client.connect()) as Model; for await (const [roundNum, round] of enumerate( - this.trainer.train(trainingDataset, validationDataset), + this.trainer.train(trainingDataset, validationDataset_), )) { yield async function* (this: Disco) { const [roundGen, roundLogsPromise] = split(round); @@ -297,6 +302,31 @@ export class Disco extends EventEmitter<{ validation.batch(batchSize).cached(), ]; } + + async #preprocessDatasets( + trainingDataset: Dataset, + validationDataset: Dataset, + ): Promise< + [ + Dataset>, + Dataset> | undefined, + ] + > { + const { batchSize } = this.#task.trainingInformation; + + let preprocessedTraining = processing.preprocess(this.#task, trainingDataset); + let preprocessedValidation = processing.preprocess(this.#task, validationDataset); + + if (this.#preprocessOnce) { + preprocessedTraining = new Dataset(await arrayFromAsync(preprocessedTraining)); + preprocessedValidation = new Dataset(await arrayFromAsync(preprocessedValidation)); + } + + return [ + preprocessedTraining.batch(batchSize).cached(), + preprocessedValidation.batch(batchSize).cached(), + ]; + } } // Array.fromAsync not yet widely used (2024) From c93631cd476a3b8ab11dc634878b82518be1517a Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 20 Apr 2026 16:50:10 +0200 Subject: [PATCH 04/89] add working local train --- cli/src/cli.ts | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 9629f5b01..9c291474b 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -26,6 +26,7 @@ import type { UserLogFile } from "./user_log.js"; async function runUser( task: Task, + provider: TaskProvider, url: URL, data: Dataset, validationData: Dataset | undefined, @@ -38,6 +39,13 @@ async function runUser( const client = clients.getClient(trainingScheme, url, task, aggregator) const disco = new Disco(task, client, { scheme: trainingScheme }); + // For local training, load model from provider before training starts + if (trainingScheme === "local") { + console.log("Loading model for local training..."); + disco.trainer.model = await provider.getModel(); + console.log("Model loaded successfully"); + } + const dir = path.join(".", `${args.testID}`); await fs.mkdir(dir, { recursive: true }); const streamPath = path.join(dir, `client${userIndex}_local_log.jsonl`); @@ -116,7 +124,7 @@ async function main( } const logs = await Promise.all( - dataSplits.map((data, i) => runUser(task, args.host, data as Dataset, validationData, i, numberOfUsers)) + dataSplits.map((data, i) => runUser(task, provider, args.host, data as Dataset, validationData, i, numberOfUsers)) ) if (args.save) { From 2e34a3737a5d71006d014d672629583ce0dcb67f Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 22 Apr 2026 16:13:52 +0200 Subject: [PATCH 05/89] add model saving to disk arg and more debug lines --- cli/src/args.ts | 2 ++ cli/src/cli.ts | 22 ++++++++++++++++++---- discojs/src/training/trainer.ts | 9 +++++++++ 3 files changed, 29 insertions(+), 4 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index 955dcf47e..00d72102e 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -38,6 +38,7 @@ export interface BenchmarkArguments { maxShareValue?: number save: boolean + saveModel: boolean host: URL } @@ -62,6 +63,7 @@ const unsafeArgs = parse( 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 }, save: { type: Boolean, alias: 's', description: 'Save logs of benchmark', defaultValue: false }, + saveModel: { type: Boolean, alias: 'm', description: 'Save trained model to disk', defaultValue: false }, host: { type: (raw: string) => new URL(raw), typeLabel: "URL", diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 9c291474b..3caf6d5cf 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -5,7 +5,7 @@ import { List, Range } from 'immutable' import fs from 'node:fs/promises' import { createWriteStream } from "node:fs"; import path from "node:path"; - +import createDebug from "debug"; import type { Dataset, DataFormat, @@ -17,12 +17,13 @@ import type { } from "@epfml/discojs"; import { Disco, aggregator as aggregators, client as clients } from '@epfml/discojs' -import { loadText } from "@epfml/discojs-node"; +import { loadText, saveModelToDisk } from "@epfml/discojs-node"; import { getTaskData } from './data.js' import { args } from './args.js' import { makeUserLogFile } from "./user_log.js"; import type { UserLogFile } from "./user_log.js"; +const debug = createDebug("cli:main"); async function runUser( task: Task, @@ -33,7 +34,8 @@ async function runUser( userIndex: number, numberOfUsers: number, ): Promise> { - // cast as typescript isn't good with generics + debug(`Starting runUser for client ${userIndex}`); + const userStart = Date.now(); const trainingScheme = task.trainingInformation.scheme as N const aggregator = aggregators.getAggregator(task) const client = clients.getClient(trainingScheme, url, task, aggregator) @@ -41,9 +43,12 @@ async function runUser( // For local training, load model from provider before training starts if (trainingScheme === "local") { + debug(`Loading model for local training client ${userIndex}...`); + const modelStart = Date.now(); console.log("Loading model for local training..."); disco.trainer.model = await provider.getModel(); console.log("Model loaded successfully"); + debug(`Model loading took ${Date.now() - modelStart}ms for client ${userIndex}`); } const dir = path.join(".", `${args.testID}`); @@ -59,6 +64,8 @@ async function runUser( } try{ + debug(`Starting training for client ${userIndex}`); + const trainStart = Date.now(); for await (const log of disco.trainSummary(data, validationData)){ finalLog.push(log); @@ -66,9 +73,16 @@ async function runUser( jsonStream.write(JSON.stringify(log) + "\n"); } } + debug(`Training took ${Date.now() - trainStart}ms for client ${userIndex}`); await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish - + // Save the trained model if requested + if (args.saveModel) { + const modelDir = path.join(".", `${args.testID}`, "models"); + const modelFileName = `client${userIndex}_model.json`; + await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); + console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); + } // saving the entire per-user logs if (args.save) { const finalPath = path.join(dir, `client${userIndex}_local_log.json`); diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 68e716bcc..73cae1f7c 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -16,8 +16,11 @@ import { } from "../index.js"; import { privacy } from "../index.js"; import { Client } from "../client/index.js"; +import createDebug from "debug"; import * as async_iterator from "../utils/async_iterator.js"; +const debug = createDebug("discojs:training:trainer"); + export interface RoundLogs { epochs: List; participants: number; @@ -88,6 +91,7 @@ export class Trainer { AsyncGenerator, RoundLogs>, void > { + debug("Start train") if (this.#training !== undefined) throw new Error( "training already running, stop it before launching a new one", @@ -109,6 +113,9 @@ export class Trainer { void > { const totalRound = Math.trunc(this.#epochs / this.#roundDuration); + + debug("Run rounds") + for (let round = 0; round < totalRound; round++) { await this.#client.onRoundBeginCommunication(); @@ -150,6 +157,8 @@ export class Trainer { ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); + debug("Run round") + // Before starting the training, get the validation of global model const validation = validationDataset !== undefined ? await this.model.evaluate(validationDataset) : undefined; From 4c9c84cc14c632e7a77b50af9b05c363f0584158 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 24 Apr 2026 17:55:42 +0200 Subject: [PATCH 06/89] add debug commands --- cli/src/cli.ts | 20 +++++++++++--------- discojs/src/client/client.ts | 5 ++++- discojs/src/default_tasks/privacyrun.ts | 6 +++--- discojs/src/models/gpt/index.ts | 8 ++++++++ discojs/src/models/gpt/model.ts | 6 ++++++ discojs/src/serialization/model.ts | 23 +++++++++++++++++++++++ discojs/src/training/disco.ts | 11 +++++++++++ 7 files changed, 66 insertions(+), 13 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 3caf6d5cf..14523d44f 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -42,15 +42,17 @@ async function runUser( const disco = new Disco(task, client, { scheme: trainingScheme }); // For local training, load model from provider before training starts - if (trainingScheme === "local") { - debug(`Loading model for local training client ${userIndex}...`); - const modelStart = Date.now(); - console.log("Loading model for local training..."); - disco.trainer.model = await provider.getModel(); - console.log("Model loaded successfully"); - debug(`Model loading took ${Date.now() - modelStart}ms for client ${userIndex}`); - } - + // if (trainingScheme === "local") { + // debug(`Loading model for training client ${userIndex}...`); + // const modelStart = Date.now(); + // console.log("Loading model for local training..."); + // disco.trainer.model = await provider.getModel(); + // console.log("Model loaded successfully"); + // debug(`Model loading took ${Date.now() - modelStart}ms for client ${userIndex}`); + // } + + + const dir = path.join(".", `${args.testID}`); await fs.mkdir(dir, { recursive: true }); const streamPath = path.join(dir, `client${userIndex}_local_log.jsonl`); diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 9e1298b93..e2168afa3 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -186,8 +186,11 @@ export abstract class Client extends EventEmitter<{ } url.pathname += `tasks/${this.task.id}/model.json` + debug("fetching latest model from server at {} for task {}...", url.href, this.task.id) + const response = await fetch(url); - if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`); + if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`) + else debug("response ok, decoding model...") const encoded = new Uint8Array(await response.arrayBuffer()) return await serialization.model.decode(encoded) diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index 9b54af387..d667e3c9c 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -1,7 +1,7 @@ import type { TaskProvider } from "../index.js"; import { Tokenizer, models, serialization } from "../index.js"; -export const privacyrun: TaskProvider<"text", "local"> = { +export const privacyrun: TaskProvider<"text", "federated"> = { async getTask() { return { id: 'privacyrun', @@ -23,9 +23,9 @@ export const privacyrun: TaskProvider<"text", "local"> = { "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", }, trainingInformation: { - scheme: 'local', + scheme: 'federated', aggregationStrategy: 'mean', - // minNbOfParticipants: 2, + minNbOfParticipants: 2, epochs: 6, validationSplit: 0.1, roundDuration: 2, diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 228d60cc4..886421892 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -228,8 +228,16 @@ export class GPT extends Model<"text"> { } static deserialize(data: GPTSerialization): Model<"text"> { + + debug("GPT model deserialization started") + const model = new GPT(data.config); + + debug("GPT model config initialized: %O", data.config) + model.weights = data.weights; + + debug("GPT model weights initialized") return model; } diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 01ee51e92..9207e8d47 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -64,8 +64,12 @@ export class GPTModel extends tf.LayersModel { let accuracyFraction: [number, number] = [0, 0]; let averageLoss = 0 let iteration = 1 + + debug("before iterator init") const iterator = await dataset.iterator() + debug("after getting iterator, before next") let next = await iterator.next() + debug("after next of iterator") while (next.done !== true && iteration <= this.config.maxIter) { let weightUpdateTime = performance.now() @@ -73,7 +77,9 @@ export class GPTModel extends tf.LayersModel { const { xs, ys } = next.value as { xs: tf.Tensor2D, ys: tf.Tensor3D } let preprocessingTime = performance.now() + debug("await batch data before {} iteration", iteration) await Promise.all([xs.data(), ys.data()]) + debug("after await batch data {} iteration", iteration) preprocessingTime = performance.now() - preprocessingTime // TODO include as a tensor inside the model diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 020d147af..4df0fab0c 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -7,6 +7,10 @@ import { GPTConfig } from '../models/index.js' import * as coder from "./coder.js"; import { Encoded, isEncoded } from "./coder.js"; +import createDebug from "debug" + +const debug = createDebug("discojs:serialization:model"); + const Type = { TFJS: 0, GPT: 1 @@ -16,11 +20,13 @@ export async function encode(model: Model): Promise { switch (true) { case model instanceof models.TFJS: { const serialized = await model.serialize(); + debug("TFJS model serialized"); return coder.encode([Type.TFJS, ...serialized]); } case model instanceof models.GPT: { const { weights, config } = model.serialize(); const serializedWeights = await serialization.weights.encode(weights); + debug("GPT model weights serialized"); return coder.encode([Type.GPT, serializedWeights, config]); } default: @@ -30,23 +36,34 @@ export async function encode(model: Model): Promise { export async function decode(encoded: Encoded): Promise> { const raw = coder.decode(encoded) + + debug("IMPORTANT:model decoded") if (!Array.isArray(raw) || raw.length < 2) { throw new Error("invalid encoding, encoding isn't an array or doesn't contain enough values") } + + debug("model encoding array length: %d", raw.length) + const type = raw[0] as unknown if (typeof type !== 'number') { throw new Error('invalid encoding, first encoding field should be the model type') } + + debug("model type: %d", type) + const rawModel = raw[1] as unknown switch (type) { case Type.TFJS: { + debug("TFJS model decoding started"); if (raw.length !== 3) throw new Error( "invalid TFJS model encoding: should be an array of length 3", ); const [rawDatatype, rawModel] = raw.slice(1) as unknown[]; + debug("TFJS model datatype: %s", rawDatatype); + let datatype; switch (rawDatatype) { case "image": @@ -70,6 +87,7 @@ export async function decode(encoded: Encoded): Promise> { if (raw.length == 2) { config = undefined } else if (raw.length == 3) { + debug("GPT model config decoding") config = raw[2] as GPTConfig } else { throw new Error('invalid encoding, gpt-tfjs model encoding should be an array of length 2 or 3') @@ -79,7 +97,12 @@ export async function decode(encoded: Encoded): Promise> { throw new Error( "invalid encoding, gpt-tfjs model weights should be an encoding of its weights", ); + + debug("GPT model weights decoding...") const weights = serialization.weights.decode(rawModel) + + debug("GPT model weights decoded, deserializing model... CONFIG MIGHT BE WRONG") + debug("GPT model config: %O", config || "undefined, using default config") return models.GPT.deserialize({weights, config}) } default: diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 7d85ba8fb..33c63c244 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -21,8 +21,12 @@ import { getAggregator } from "../aggregator/index.js"; import { enumerate, split } from "../utils/async_iterator.js"; import { EventEmitter } from "../utils/event_emitter.js"; +import createDebug from "debug" + import { RoundLogs, Trainer } from "./trainer.js"; +const debug = createDebug("discojs:training:disco"); + interface DiscoConfig { scheme: N; logger: Logger; @@ -175,6 +179,8 @@ export class Disco extends EventEmitter<{ const epochResults: Array<{epochNum: number; epochLogs: EpochLogs}> = []; + debug("Starting round %d", roundNum) + for await (const [epochNum, epoch] of enumerate(roundGen)) { const [epochGen, epochLogsPromise] = async_iterator.split(epoch); for await (const _ of epochGen); @@ -218,7 +224,12 @@ export class Disco extends EventEmitter<{ // the client fetches the latest weights upon connection // TODO unsafe cast + debug("Connecting to client and fetching initial model..."); this.trainer.model = (await this.#client.connect()) as Model; + debug("Initial model fetched successfully"); + if (this.trainer.model === null) { + debug(`No pre-trained model provided for client, initializing randomly...`); + } for await (const [roundNum, round] of enumerate( this.trainer.train(trainingDataset, validationDataset_), From 23c733d5a4817b81e17c9a51da782b402d128ee9 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 24 Apr 2026 18:21:56 +0200 Subject: [PATCH 07/89] add cnahges to federated approach --- discojs/src/client/client.ts | 2 +- discojs/src/client/event_connection.ts | 4 ++++ discojs/src/client/federated/federated_client.ts | 4 +++- discojs/src/client/federated/messages.ts | 2 +- server/src/controllers/federated_controller.ts | 2 +- 5 files changed, 10 insertions(+), 4 deletions(-) diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index e2168afa3..5c8aa1304 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -186,7 +186,7 @@ export abstract class Client extends EventEmitter<{ } url.pathname += `tasks/${this.task.id}/model.json` - debug("fetching latest model from server at {} for task {}...", url.href, this.task.id) + debug("fetching latest model from server at %0 for task %1...", url.href, this.task.id) const response = await fetch(url); if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`) diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index 3e3aec409..82722d111 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -118,6 +118,10 @@ export class WebSocketServer extends EventEmitter<{ [K in type]: NarrowMessage { + debug("websocket closed: code=%o reason=%o wasClean=%o", event.code, event.reason, event.wasClean) + } + return await new Promise((resolve, reject) => { ws.onerror = (err: WebSocket.ErrorEvent) => { reject(new Error(`Server unreachable: ${err.message}`)) diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index a89c65ad6..daaa1df80 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -88,7 +88,9 @@ export class FederatedClient extends Client<"federated"> { // Upon connecting, the server answers with a boolean // which indicates whether there are enough participants or not debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants) - model.weights = serialization.weights.decode(payload) + if (payload !== undefined) { + model.weights = serialization.weights.decode(payload) + } return model } diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts index 3733d2c1c..d8961d97d 100644 --- a/discojs/src/client/federated/messages.ts +++ b/discojs/src/client/federated/messages.ts @@ -18,7 +18,7 @@ export interface NewFederatedNodeInfo { type: type.NewFederatedNodeInfo id: NodeID waitForMoreParticipants: boolean - payload: serialization.Encoded; + payload?: serialization.Encoded; round: number nbOfParticipants: number } diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 1e52fe539..cf2df0fa8 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -89,7 +89,7 @@ export class FederatedController extends TrainingController< type: MessageTypes.NewFederatedNodeInfo, id: clientId, waitForMoreParticipants: this.connections.size < minNbOfParticipants, - payload: this.#latestGlobalWeights, + payload: this.#aggregator.round === 0 ? undefined : this.#latestGlobalWeights, round: this.#aggregator.round, nbOfParticipants: this.connections.size } From 096393c8b70d57348e99f6bc5ecbd966d5804c26 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 24 Apr 2026 18:31:27 +0200 Subject: [PATCH 08/89] change round 0 payload null handling --- discojs/src/client/federated/federated_client.ts | 2 +- discojs/src/client/federated/messages.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index daaa1df80..b6c2c59d5 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -88,7 +88,7 @@ export class FederatedClient extends Client<"federated"> { // Upon connecting, the server answers with a boolean // which indicates whether there are enough participants or not debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants) - if (payload !== undefined) { + if (payload != null) { model.weights = serialization.weights.decode(payload) } return model diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts index d8961d97d..c6dee6698 100644 --- a/discojs/src/client/federated/messages.ts +++ b/discojs/src/client/federated/messages.ts @@ -18,7 +18,7 @@ export interface NewFederatedNodeInfo { type: type.NewFederatedNodeInfo id: NodeID waitForMoreParticipants: boolean - payload?: serialization.Encoded; + payload?: serialization.Encoded | null; round: number nbOfParticipants: number } From 1ef7d856b18a5a1a08f59e44336623c82b82d3f3 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 26 Apr 2026 22:51:57 +0200 Subject: [PATCH 09/89] chnage server max payload limit to higher number --- server/src/controllers/federated_controller.ts | 4 ++++ server/src/server.ts | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index cf2df0fa8..9da1e589a 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -66,6 +66,10 @@ export class FederatedController extends TrainingController< } const shortId = clientId.slice(0, 4) + ws.on('error', (err) => { + debug("websocket error for client [%s]: %o", shortId, err) + }) + // Setup callbacks triggered upon receiving the different client messages ws.on('message', (data: Buffer) => { const msg: unknown = msgpack.decode(data) diff --git a/server/src/server.ts b/server/src/server.ts index 06641cff0..b862c8272 100644 --- a/server/src/server.ts +++ b/server/src/server.ts @@ -38,6 +38,10 @@ export class Server { async serve(port?: number): Promise<[http.Server, URL]> { const wsApplier = expressWS(express(), undefined, { leaveRouterUntouched: true, + wsOptions: { + // GPT-sized federated updates can exceed the ws default payload limit. + maxPayload: 1024 * 1024 * 1024, + }, }); const app = wsApplier.app; From 417bfa5967d64e6ef0121dcea471775e75bfd97c Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 26 Apr 2026 23:03:19 +0200 Subject: [PATCH 10/89] fix memory reads and wait for all clients to begin --- cli/src/cli.ts | 2 +- discojs/src/client/federated/federated_client.ts | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 14523d44f..51697cde1 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -39,7 +39,7 @@ async function runUser( const trainingScheme = task.trainingInformation.scheme as N const aggregator = aggregators.getAggregator(task) const client = clients.getClient(trainingScheme, url, task, aggregator) - const disco = new Disco(task, client, { scheme: trainingScheme }); + const disco = new Disco(task, client, { scheme: trainingScheme, preprocessOnce: true }); // For local training, load model from provider before training starts // if (trainingScheme === "local") { diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index b6c2c59d5..9caea57f0 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -105,9 +105,16 @@ export class FederatedClient extends Client<"federated"> { this.aggregator.setNodes(this.aggregator.nodes.delete(SERVER_NODE_ID)); } - override onRoundBeginCommunication(): Promise { + // override onRoundBeginCommunication(): Promise { + // // Prepare the result promise for the incoming round + // this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve)) + // this.saveAndEmit("local training") + // return Promise.resolve(); + // } + override async onRoundBeginCommunication(): Promise { // Prepare the result promise for the incoming round this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve)) + await this.waitForParticipantsIfNeeded() // In case we are waiting for more participants, we wait before starting the local training this.saveAndEmit("local training") return Promise.resolve(); } From 77ce9a9cf24c0ade32ceb346d81570e85d3f3ab5 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 27 Apr 2026 01:28:54 +0200 Subject: [PATCH 11/89] add server debug logs to see why ws close session --- discojs/src/client/event_connection.ts | 6 +++++- server/src/controllers/federated_controller.ts | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index 82722d111..a8c562a74 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -98,7 +98,10 @@ export class WebSocketServer extends EventEmitter<{ [K in type]: NarrowMessage msg is Message, validateSent: (msg: Message) => boolean): Promise { - const ws = new WebSocket(url) + const ws = new WebSocket(url, { + // Federated GPT updates can exceed the default ws payload limit. + maxPayload: 1024 * 1024 * 1024, + }) ws.binaryType = 'arraybuffer' const server: WebSocketServer = new WebSocketServer(ws, validateSent) @@ -124,6 +127,7 @@ export class WebSocketServer extends EventEmitter<{ [K in type]: NarrowMessage { ws.onerror = (err: WebSocket.ErrorEvent) => { + debug("websocket error while connecting/receiving: %o", err.message) reject(new Error(`Server unreachable: ${err.message}`)) } ws.onopen = () => { resolve(server) } diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 9da1e589a..bef9a1909 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -108,6 +108,12 @@ export class FederatedController extends TrainingController< case MessageTypes.SendPayload: { const { payload, round } = msg if (this.#aggregator.isValidContribution(clientId, round)) { + debug( + "Received valid contribution from client [%s] for round %d (participants=%d)", + shortId, + round, + this.connections.size, + ) const weights = serialization.weights.decode(payload) // Create a callback to send the aggregated weight to the client @@ -120,9 +126,12 @@ export class FederatedController extends TrainingController< payload: await serialization.weights.encode(weightUpdate), nbOfParticipants: this.connections.size } + debug("Prepared aggregated payload for client [%s] at round %o", shortId, this.#aggregator.round) ws.send(msgpack.encode(msg)) + debug("Aggregated payload sent to client [%s] for round %o", shortId, this.#aggregator.round) }) // Add the contribution + debug("Adding contribution from client [%s] to aggregator for round %d", shortId, round) this.#aggregator.add(clientId, weights, round) debug(`Successfully added contribution from client [%s] for round ${round}`, shortId) } else { From 7aef628633d7c1220060d58efa1f42598f713b70 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 27 Apr 2026 02:14:56 +0200 Subject: [PATCH 12/89] cover whole dataset and split data to clients --- cli/src/data.ts | 58 ++++++++++++++++++++++++++++++-- discojs/src/models/gpt/config.ts | 3 +- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/cli/src/data.ts b/cli/src/data.ts index f3b6ff834..0f3d22e62 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -1,4 +1,5 @@ import path from "node:path"; +import { createReadStream } from "node:fs"; import { Dataset, processing } from "@epfml/discojs"; import { DataFormat, @@ -10,6 +11,44 @@ import { import { loadCSV, loadImage, loadImagesInDir, loadText } from "@epfml/discojs-node"; import { Repeat } from "immutable"; +function loadShardedTextSamples( + filePath: string, + userIdx: number, + totalClient: number, +): Dataset { + return new Dataset(async function* () { + const stream = createReadStream(filePath, { encoding: "utf8" }); + const sampleDelimiter = "<|endoftext|>"; + let buffer = ""; + let sampleIndex = 0; + + for await (const chunk of stream) { + if (typeof chunk !== "string") { + throw new Error("Expected file stream to yield string"); + } + + buffer += chunk; + + let delimiterIndex = buffer.indexOf(sampleDelimiter); + while (delimiterIndex !== -1) { + const sample = buffer.slice(0, delimiterIndex + sampleDelimiter.length).trim(); + if (sample !== "" && sampleIndex % totalClient === userIdx) { + yield sample; + } + + sampleIndex++; + buffer = buffer.slice(delimiterIndex + sampleDelimiter.length); + delimiterIndex = buffer.indexOf(sampleDelimiter); + } + } + + const trailingSample = buffer.trim(); + if (trailingSample !== "" && sampleIndex % totalClient === userIdx) { + yield trailingSample; + } + }); +} + async function loadSimpleFaceData(userIdx: number, totalClient: number): Promise> { const folder = path.join("..", "datasets", "simple_face"); @@ -122,8 +161,23 @@ export async function getTaskData( case "mnist_federated": case "mnist": return loadData("mnist", userIdx) as Dataset; - case "privacyrun": - return loadText(isValidation && validationDatasetPath ? validationDatasetPath : datasetPath ?? '../datasets/med_mcq/train.txt') as Dataset; + case "privacyrun": { + const filePath = + isValidation && validationDatasetPath + ? validationDatasetPath + : datasetPath ?? "../datasets/med_mcq/train.txt"; + + // Keep validation shared, but shard training data across clients by MCQ sample. + if (isValidation) { + return loadText(filePath) as Dataset; + } + + return loadShardedTextSamples( + filePath, + userIdx, + totalClient, + ) as Dataset; + } default: throw new Error(`Data loader for ${taskID} not implemented.`); } diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index 0d157a7fd..b2a86340f 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -31,7 +31,8 @@ export type GPTConfig = { export const DefaultGPTConfig: Required = { lr: 0.001, weightDecay: 0, - maxIter: 10, + // By default, iterate through the whole dataset and let dataset exhaustion stop the epoch. + maxIter: Number.MAX_SAFE_INTEGER, verbose: 0, modelType: 'gpt-nano', evaluate: true, From d65fb4a6692977a6063bfb3dbc0d7bbc8653d271 Mon Sep 17 00:00:00 2001 From: Mina Petrovic Date: Mon, 4 May 2026 14:58:28 +0200 Subject: [PATCH 13/89] add validation dataset loading changes --- cli/src/cli.ts | 7 ++++--- cli/src/data.ts | 26 ++++++++++++++++++-------- datasets/.gitignore | 1 + 3 files changed, 23 insertions(+), 11 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 51697cde1..5540974ae 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -17,7 +17,7 @@ import type { } from "@epfml/discojs"; import { Disco, aggregator as aggregators, client as clients } from '@epfml/discojs' -import { loadText, saveModelToDisk } from "@epfml/discojs-node"; +import { saveModelToDisk } from "@epfml/discojs-node"; import { getTaskData } from './data.js' import { args } from './args.js' import { makeUserLogFile } from "./user_log.js"; @@ -135,8 +135,9 @@ async function main( let validationData: Dataset | undefined = undefined; if (args.validationDatasetPath) { - // Assume text task for now - validationData = loadText(args.validationDatasetPath).cached() as Dataset; + validationData = ( + await getTaskData(task.id, 0, 1, args.validationDatasetPath, true, args.validationDatasetPath) + ).cached() as Dataset; } const logs = await Promise.all( diff --git a/cli/src/data.ts b/cli/src/data.ts index 0f3d22e62..a1f41e9aa 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -8,13 +8,13 @@ import { Task, Text, } from "@epfml/discojs"; -import { loadCSV, loadImage, loadImagesInDir, loadText } from "@epfml/discojs-node"; +import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node"; import { Repeat } from "immutable"; -function loadShardedTextSamples( +function loadTextSamples( filePath: string, - userIdx: number, - totalClient: number, + userIdx?: number, + totalClient?: number, ): Dataset { return new Dataset(async function* () { const stream = createReadStream(filePath, { encoding: "utf8" }); @@ -32,7 +32,12 @@ function loadShardedTextSamples( let delimiterIndex = buffer.indexOf(sampleDelimiter); while (delimiterIndex !== -1) { const sample = buffer.slice(0, delimiterIndex + sampleDelimiter.length).trim(); - if (sample !== "" && sampleIndex % totalClient === userIdx) { + const shouldYield = + userIdx === undefined || + totalClient === undefined || + sampleIndex % totalClient === userIdx; + + if (sample !== "" && shouldYield) { yield sample; } @@ -43,7 +48,12 @@ function loadShardedTextSamples( } const trailingSample = buffer.trim(); - if (trailingSample !== "" && sampleIndex % totalClient === userIdx) { + const shouldYieldTrailing = + userIdx === undefined || + totalClient === undefined || + sampleIndex % totalClient === userIdx; + + if (trailingSample !== "" && shouldYieldTrailing) { yield trailingSample; } }); @@ -169,10 +179,10 @@ export async function getTaskData( // Keep validation shared, but shard training data across clients by MCQ sample. if (isValidation) { - return loadText(filePath) as Dataset; + return loadTextSamples(filePath) as Dataset; } - return loadShardedTextSamples( + return loadTextSamples( filePath, userIdx, totalClient, diff --git a/datasets/.gitignore b/datasets/.gitignore index e644c626a..bdbf3ed6d 100644 --- a/datasets/.gitignore +++ b/datasets/.gitignore @@ -7,6 +7,7 @@ /simple_face-example.png /titanic* /mnist* +/medicalMCQtxtNoExplanation # wikitext /wikitext/ From fe7c51a98eb8e8c707a2703a98af5f887789713b Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 4 May 2026 15:05:44 +0200 Subject: [PATCH 14/89] change gpt config to use whole dadataset --- discojs/src/models/gpt/index.ts | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 886421892..082f79766 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -231,9 +231,17 @@ export class GPT extends Model<"text"> { debug("GPT model deserialization started") - const model = new GPT(data.config); + const config = + data.config === undefined + ? undefined + : { + ...data.config, + maxIter: DefaultGPTConfig.maxIter, + }; - debug("GPT model config initialized: %O", data.config) + const model = new GPT(config); + + debug("GPT model config initialized: %O", config) model.weights = data.weights; From 0dc32aabf4fcda8595a9fc816bd728e183200c68 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 5 May 2026 17:51:51 +0200 Subject: [PATCH 15/89] add arg for model saving location, cnahge save to saveLog, change link for model to contxt 512 --- cli/src/args.ts | 6 ++++-- cli/src/cli.ts | 18 +++++++++++------- discojs/src/default_tasks/privacyrun.ts | 11 +++++++---- 3 files changed, 22 insertions(+), 13 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index 00d72102e..984db0156 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -23,6 +23,7 @@ export interface BenchmarkArguments { validationSplit: number datasetPath?: string validationDatasetPath?: string + outputPath?: string // DP epsilon?: number @@ -37,7 +38,7 @@ export interface BenchmarkArguments { // Secure aggregator maxShareValue?: number - save: boolean + saveLogs: boolean saveModel: boolean host: URL } @@ -62,7 +63,8 @@ const unsafeArgs = parse( validationSplit : { type: Number, alias: 'v', description: 'Validation dataset ratio', defaultValue: 0.2 }, 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 }, - save: { type: Boolean, alias: 's', description: 'Save logs of benchmark', defaultValue: false }, + outputPath: { type: String, alias: 'o', description: 'Path to save logs and models. Defaults to ./', 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 }, host: { type: (raw: string) => new URL(raw), diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 5540974ae..e780f6b65 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -25,6 +25,10 @@ import type { UserLogFile } from "./user_log.js"; const debug = createDebug("cli:main"); +function getOutputDir(): string { + return args.outputPath ?? path.join(".", `${args.testID}`); +} + async function runUser( task: Task, provider: TaskProvider, @@ -39,7 +43,7 @@ async function runUser( const trainingScheme = task.trainingInformation.scheme as N const aggregator = aggregators.getAggregator(task) const client = clients.getClient(trainingScheme, url, task, aggregator) - const disco = new Disco(task, client, { scheme: trainingScheme, preprocessOnce: true }); + const disco = new Disco(task, client, { scheme: trainingScheme, preprocessOnce: false }); // For local training, load model from provider before training starts // if (trainingScheme === "local") { @@ -53,7 +57,7 @@ async function runUser( - const dir = path.join(".", `${args.testID}`); + const dir = getOutputDir(); await fs.mkdir(dir, { recursive: true }); const streamPath = path.join(dir, `client${userIndex}_local_log.jsonl`); @@ -61,7 +65,7 @@ async function runUser( // create a write stream that saves learning logs during the train let jsonStream: ReturnType | null = null; - if (args.save){ + if (args.saveLogs){ jsonStream = createWriteStream(streamPath, {flags: "w"}); } @@ -80,13 +84,13 @@ async function runUser( await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish // Save the trained model if requested if (args.saveModel) { - const modelDir = path.join(".", `${args.testID}`, "models"); + const modelDir = path.join(getOutputDir(), "models"); const modelFileName = `client${userIndex}_model.json`; await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); } // saving the entire per-user logs - if (args.save) { + if (args.saveLogs) { const finalPath = path.join(dir, `client${userIndex}_local_log.json`); const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, client.ownId, finalLog); @@ -144,8 +148,8 @@ async function main( dataSplits.map((data, i) => runUser(task, provider, args.host, data as Dataset, validationData, i, numberOfUsers)) ) - if (args.save) { - const dir = path.join(".", `${args.testID}`, `${task.id}`); + if (args.saveLogs) { + const dir = path.join(getOutputDir(), `${task.id}`); await fs.mkdir(dir, { recursive: true }); const filePath = path.join(dir, `${task.id}_${numberOfUsers}users.json`); diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index d667e3c9c..dcdaa7640 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -26,12 +26,13 @@ export const privacyrun: TaskProvider<"text", "federated"> = { scheme: 'federated', aggregationStrategy: 'mean', minNbOfParticipants: 2, - epochs: 6, + epochs: 1, validationSplit: 0.1, - roundDuration: 2, + roundDuration: 1, batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), - contextLength: 1024, + // contextLength: 1024, + contextLength: 512, tensorBackend: 'gpt' } } @@ -40,8 +41,10 @@ export const privacyrun: TaskProvider<"text", "federated"> = { async getModel() { // Load the pre-trained ONNX-converted model from Google Cloud Storage // The model should be in DiscoJS serialization format (created by onnx-converter) - const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model.json"; + // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model.json"; + const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; + try { const response = await fetch(modelUrl); if (!response.ok) { From 5c03f40a98bde019ffda88b96a70b8556a939680 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 5 May 2026 18:02:42 +0200 Subject: [PATCH 16/89] add training optimizations --- discojs/src/models/gpt/model.ts | 36 +++++++++++++++++---------------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 9207e8d47..a69e498b4 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -78,27 +78,29 @@ export class GPTModel extends tf.LayersModel { let preprocessingTime = performance.now() debug("await batch data before {} iteration", iteration) - await Promise.all([xs.data(), ys.data()]) + // await Promise.all([xs.data(), ys.data()]) + await Promise.resolve() debug("after await batch data {} iteration", iteration) preprocessingTime = performance.now() - preprocessingTime // TODO include as a tensor inside the model - const accTensor = tf.tidy(() => { - const logits = this.apply(xs) - if (Array.isArray(logits)) - throw new Error('model outputs too many tensor') - if (logits instanceof tf.SymbolicTensor) - throw new Error('model outputs symbolic tensor') - return tf.metrics.categoricalAccuracy(ys, logits) - }) - const accSize = accTensor.shape.reduce((l, r) => l * r, 1) - const accSumTensor = accTensor.sum() - const accSum = await accSumTensor.array() - tf.dispose(accSumTensor) - if (typeof accSum !== 'number') - throw new Error('got multiple accuracy sum') - accuracyFraction = [accuracyFraction[0] + accSum, accuracyFraction[1] + accSize]; - tf.dispose([accTensor]) + // const accTensor = tf.tidy(() => { + // const logits = this.apply(xs) + // if (Array.isArray(logits)) + // throw new Error('model outputs too many tensor') + // if (logits instanceof tf.SymbolicTensor) + // throw new Error('model outputs symbolic tensor') + // return tf.metrics.categoricalAccuracy(ys, logits) + // }) + // const accSize = accTensor.shape.reduce((l, r) => l * r, 1) + // const accSumTensor = accTensor.sum() + // const accSum = await accSumTensor.array() + // tf.dispose(accSumTensor) + // if (typeof accSum !== 'number') + // throw new Error('got multiple accuracy sum') + // accuracyFraction = [accuracyFraction[0] + accSum, accuracyFraction[1] + accSize]; + // tf.dispose([accTensor]) + accuracyFraction = [Number.NaN, Number.NaN]; const lossTensor = tf.tidy(() => { const { grads, value: lossTensor } = this.optimizer.computeGradients(() => { From 4ee6096b4ba35e3621af3a29c82323ccf8d5c6df Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 6 May 2026 19:13:21 +0200 Subject: [PATCH 17/89] change onnx converter to be able to convert different context len --- onnx-converter/src/convert_onnx.ts | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/onnx-converter/src/convert_onnx.ts b/onnx-converter/src/convert_onnx.ts index cf45a76ed..dd3d1aab2 100644 --- a/onnx-converter/src/convert_onnx.ts +++ b/onnx-converter/src/convert_onnx.ts @@ -7,6 +7,7 @@ import { models, serialization } from "@epfml/discojs"; const OUTPUT_FILENAME = "model.json"; const GPT2_N_LAYER = 12; +const GPT2_CONTEXT_LENGTH = 1024; const ONNX_URL = "https://huggingface.co/Xenova/gpt2/resolve/main/onnx/decoder_model.onnx?download=true" @@ -29,8 +30,7 @@ async function main() { // Init empty TF.js model - // Context length value from https://huggingface.co/Xenova/gpt2/blob/main/config.json - const gptModel = new models.GPT({ modelType: 'gpt2', contextLength: 1024 }); + const gptModel = new models.GPT({ modelType: 'gpt2', contextLength: GPT2_CONTEXT_LENGTH }); if (gptModel.config.nLayer != GPT2_N_LAYER) throw new Error(`ONNX conversion only supports GPT-2 with 12 layers, instead found ${gptModel.config.nLayer}.`); const gptLayersModel = gptModel.extract(); @@ -54,7 +54,14 @@ async function main() { throw new Error(`Undefined layer dimensions for ${tensor.name}`) const dims = tensor.dims.map((d) => Number(d)); const flatData = parseTensorData(tensor); - const tfTensor = tf.tensor(flatData).reshape(dims) + let tfTensor = tf.tensor(flatData).reshape(dims) + if (tensor.name === "transformer.wpe.weight") { + if (dims.length !== 2) + throw new Error(`Expected transformer.wpe.weight to be a 2D tensor, got ${dims.length}D.`); + if (dims[0] < GPT2_CONTEXT_LENGTH) + throw new Error(`ONNX positional embeddings only support context length ${dims[0]}, requested ${GPT2_CONTEXT_LENGTH}.`); + tfTensor = tfTensor.slice([0, 0], [GPT2_CONTEXT_LENGTH, dims[1]]); + } preTrainedWeights = preTrainedWeights.set(tfjsName, tfTensor); } @@ -133,4 +140,4 @@ function createWeightNameMap(): Map { } -await main().catch(console.error); \ No newline at end of file +await main().catch(console.error); From 5b5f88cc892ad582a9dd3455c27b1daa8e82f950 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 11 May 2026 15:06:08 +0200 Subject: [PATCH 18/89] aggregate inside of an epoch for llms --- cli/package.json | 1 + cli/src/args.ts | 3 + cli/src/measure_memorization_gpt2.ts | 316 +++++++++++++++++++++++ discojs/src/models/gpt/index.ts | 36 ++- discojs/src/models/gpt/model.ts | 16 +- discojs/src/task/training_information.ts | 2 + discojs/src/training/trainer.ts | 133 +++++++++- 7 files changed, 499 insertions(+), 8 deletions(-) create mode 100644 cli/src/measure_memorization_gpt2.ts diff --git a/cli/package.json b/cli/package.json index 8ed779b3f..6c4f10bb3 100644 --- a/cli/package.json +++ b/cli/package.json @@ -10,6 +10,7 @@ "train_gpt": "npm run build && node dist/train_gpt.js", "hellaswag_gpt": "npm run build && node dist/hellaswag_gpt.js", "eval_finetuned_gpt2": "npm run build && node dist/evaluate_finetuned_gpt2.js", + "measure_memorization_gpt2": "npm run build && node dist/measure_memorization_gpt2.js", "finetune_gpt": "npm run build && node dist/finetune_gpt.js", "build": "tsc --build", "test": ": nothing" diff --git a/cli/src/args.ts b/cli/src/args.ts index 984db0156..f7755b43d 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -19,6 +19,7 @@ export interface BenchmarkArguments { numberOfUsers: number epochs: number roundDuration: number + roundIterations?: number batchSize: number validationSplit: number datasetPath?: string @@ -59,6 +60,7 @@ const unsafeArgs = parse( numberOfUsers: { type: Number, alias: 'u', description: 'Number of users', defaultValue: 2 }, epochs: { type: Number, alias: 'e', description: 'Number of epochs', defaultValue: 10 }, roundDuration: { type: Number, alias: 'r', 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', description: 'Training batch size', defaultValue: 10 }, validationSplit : { type: Number, alias: 'v', description: 'Validation dataset ratio', defaultValue: 0.2 }, datasetPath: { type: String, alias: 'd', description: 'Path to the dataset', optional: true }, @@ -133,6 +135,7 @@ export const args: BenchmarkArguments = { task.trainingInformation.roundDuration = unsafeArgs.roundDuration; task.trainingInformation.epochs = unsafeArgs.epochs; task.trainingInformation.validationSplit = unsafeArgs.validationSplit; + (task.trainingInformation as typeof task.trainingInformation & { roundIterations?: number }).roundIterations = unsafeArgs.roundIterations; const {aggregator, clippingRadius, maxIterations, beta, maxShareValue} = unsafeArgs; diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts new file mode 100644 index 000000000..63f02d1ac --- /dev/null +++ b/cli/src/measure_memorization_gpt2.ts @@ -0,0 +1,316 @@ +import "@tensorflow/tfjs-node"; +import * as tf from "@tensorflow/tfjs"; +import fs from "node:fs/promises"; +import { parse } from "ts-command-line-args"; + +import { models, Tokenizer } from "@epfml/discojs"; +import { loadModelFromDisk } from "@epfml/discojs-node"; + +interface Args { + modelPath: string; + dataPath: string; + maxRecords: number; + promptLengths: string; + suffixLength: number; + bleuThreshold: number; + seed: number; + savePath?: string; + help?: boolean; +} + +type PromptResult = { + recordIndex: number; + promptLength: number; + splitIndex: number; + exactMatch: boolean; + bleu: number; + memorizedByBleu: boolean; + promptText: string; + referenceText: string; + generatedText: string; +}; + +function parseIntegerList(raw: string): number[] { + const values = raw + .split(",") + .map((v) => Number.parseInt(v.trim(), 10)) + .filter((v) => !Number.isNaN(v)); + + if (values.length === 0 || values.some((v) => v <= 0)) { + throw new Error("promptLengths must be a comma-separated list of positive integers"); + } + + return values; +} + +function seededRandom(seed: number): () => number { + let state = seed >>> 0; + return () => { + state = (1664525 * state + 1013904223) >>> 0; + return state / 0x100000000; + }; +} + +function randomInt(random: () => number, minInclusive: number, maxInclusive: number): number { + if (maxInclusive < minInclusive) { + throw new Error("invalid random integer range"); + } + + return minInclusive + Math.floor(random() * (maxInclusive - minInclusive + 1)); +} + +async function loadRecords(filePath: string, limit: number): Promise { + const text = await fs.readFile(filePath, "utf8"); + const delimiter = "<|endoftext|>"; + const rawRecords = text.includes(delimiter) + ? text.split(delimiter) + : text.split(/\n\s*\n/g); + + const records = rawRecords + .map((record) => + record + .replaceAll("<|startoftext|>", "") + .replaceAll("<|endoftext|>", "") + .trim(), + ) + .filter((record) => record.length > 0); + + return limit > 0 ? records.slice(0, limit) : records; +} + +function ngrams(tokens: number[], n: number): Map { + const counts = new Map(); + if (tokens.length < n) return counts; + + for (let i = 0; i <= tokens.length - n; i++) { + const key = tokens.slice(i, i + n).join(","); + counts.set(key, (counts.get(key) ?? 0) + 1); + } + + return counts; +} + +function bleu1to4(reference: number[], candidate: number[]): number { + if (candidate.length === 0) return 0; + + const precisions: number[] = []; + for (let n = 1; n <= 4; n++) { + const referenceCounts = ngrams(reference, n); + const candidateCounts = ngrams(candidate, n); + let overlap = 0; + let total = 0; + + for (const [key, count] of candidateCounts) { + overlap += Math.min(count, referenceCounts.get(key) ?? 0); + total += count; + } + + precisions.push(total === 0 ? 0 : overlap / total); + } + + if (precisions.some((precision) => precision === 0)) return 0; + + const brevityPenalty = + candidate.length > reference.length + ? 1 + : Math.exp(1 - reference.length / candidate.length); + const geometricMean = Math.exp( + precisions.reduce((sum, precision) => sum + Math.log(precision), 0) / precisions.length, + ); + + return brevityPenalty * geometricMean; +} + +async function greedyGenerateGPT2( + model: models.GPT, + inputIds: number[], + maxNewTokens: number, + maxContextLength: number, +): Promise { + const generated = [...inputIds]; + const tfModel = model.extract(); + + 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 output = tfModel.predict(input); + if (Array.isArray(output)) { + return output[0] as tf.Tensor; + } + return output as tf.Tensor; + }); + + const nextTokenTensor = tf.tidy(() => { + const last = logits.slice([0, modelInput.length - 1, 0], [1, 1, -1]); + return last.squeeze().argMax(); + }); + + const nextTokenData = await nextTokenTensor.data(); + const nextToken = nextTokenData[0]; + + input.dispose(); + logits.dispose(); + nextTokenTensor.dispose(); + + generated.push(nextToken); + } + + return generated; +} + +function summarize(results: PromptResult[]) { + const byPromptLength = new Map(); + for (const result of results) { + byPromptLength.set( + result.promptLength, + [...(byPromptLength.get(result.promptLength) ?? []), result], + ); + } + + const summarizeGroup = (group: PromptResult[]) => ({ + count: group.length, + exactMatchRate: group.filter((r) => r.exactMatch).length / group.length, + bleuMemorizationRate: group.filter((r) => r.memorizedByBleu).length / group.length, + averageBleu: group.reduce((sum, r) => sum + r.bleu, 0) / group.length, + }); + + return { + overall: summarizeGroup(results), + byPromptLength: Object.fromEntries( + [...byPromptLength.entries()].map(([promptLength, group]) => [ + promptLength, + summarizeGroup(group), + ]), + ), + }; +} + +async function main() { + const args = parse( + { + modelPath: { type: String, description: "Path to a saved Disco GPT model.json" }, + dataPath: { type: String, description: "Path to records/canaries text file" }, + maxRecords: { type: Number, description: "Maximum records to evaluate; -1 for all", defaultValue: 100 }, + promptLengths: { type: String, description: "Comma-separated prompt lengths", defaultValue: "10,50,100,200,500" }, + suffixLength: { type: Number, description: "Number of suffix tokens to generate and compare", defaultValue: 50 }, + bleuThreshold: { type: Number, description: "BLEU threshold for approximate memorization", defaultValue: 0.75 }, + seed: { type: Number, description: "Random seed for choosing record split positions", defaultValue: 42 }, + savePath: { type: String, description: "Optional JSON output path", optional: true }, + help: { type: Boolean, optional: true, alias: "h", description: "Prints this usage guide" }, + }, + { + helpArg: "help", + headerContentSections: [ + { + header: "GPT-2 Unintended Memorization", + content: "Measures extractable memorization via greedy suffix generation.", + }, + ], + }, + ); + + const promptLengths = parseIntegerList(args.promptLengths); + const maxPromptLength = Math.max(...promptLengths); + const random = seededRandom(args.seed); + + console.log("Loading tokenizer..."); + const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + + console.log("Loading model..."); + const loadedModel = await loadModelFromDisk(args.modelPath); + if (!(loadedModel instanceof models.GPT)) { + throw new Error("modelPath must point to a Disco GPT model"); + } + + console.log("Loading records..."); + const records = await loadRecords(args.dataPath, args.maxRecords); + console.log(`Loaded ${records.length} records`); + + const results: PromptResult[] = []; + let skipped = 0; + + for (let recordIndex = 0; recordIndex < records.length; recordIndex++) { + const record = records[recordIndex]; + const ids = tokenizer.tokenize(record).toArray(); + + if (ids.length < maxPromptLength + args.suffixLength + 1) { + skipped++; + continue; + } + + const splitIndex = randomInt( + random, + maxPromptLength, + ids.length - args.suffixLength, + ); + const reference = ids.slice(splitIndex, splitIndex + args.suffixLength); + + for (const promptLength of promptLengths) { + const prompt = ids.slice(splitIndex - promptLength, splitIndex); + const generated = await greedyGenerateGPT2( + loadedModel, + prompt, + args.suffixLength, + loadedModel.config.contextLength, + ); + const generatedSuffix = generated.slice(prompt.length, prompt.length + args.suffixLength); + const exactMatch = + generatedSuffix.length === reference.length && + generatedSuffix.every((token, i) => token === reference[i]); + const bleu = bleu1to4(reference, generatedSuffix); + + results.push({ + recordIndex, + promptLength, + splitIndex, + exactMatch, + bleu, + memorizedByBleu: bleu > args.bleuThreshold, + promptText: tokenizer.decode(prompt), + referenceText: tokenizer.decode(reference), + generatedText: tokenizer.decode(generatedSuffix), + }); + } + + if ((recordIndex + 1) % 10 === 0) { + console.log(`Processed ${recordIndex + 1}/${records.length} records`); + } + } + + if (results.length === 0) { + throw new Error("No records were long enough to evaluate"); + } + + const summary = { + config: { + modelPath: args.modelPath, + dataPath: args.dataPath, + maxRecords: args.maxRecords, + promptLengths, + suffixLength: args.suffixLength, + bleuThreshold: args.bleuThreshold, + seed: args.seed, + modelContextLength: loadedModel.config.contextLength, + }, + skippedRecords: skipped, + ...summarize(results), + }; + + console.log("\n=== Memorization Summary ==="); + console.log(JSON.stringify(summary, null, 2)); + + if (args.savePath !== undefined) { + await fs.writeFile( + args.savePath, + JSON.stringify({ summary, results }, null, 2), + ); + console.log(`Saved detailed results to ${args.savePath}`); + } +} + +main().catch((err) => { + console.error(err); + process.exitCode = 1; +}); diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 082f79766..458dad7e6 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -30,6 +30,7 @@ export class GPT extends Model<"text"> { readonly #contextLength: number; readonly #maxBatchCount: number; readonly #vocabSize: number; + #iterationCount = 0; constructor(partialConfig?: Partial, layersModel?: tf.LayersModel) { super(); @@ -62,7 +63,7 @@ export class GPT extends Model<"text"> { for await (const [batch, _] of trainingDataset.zip( Range(0, this.#maxBatchCount), )) { - const batchLogs = await this.#runBatch(batch); + const batchLogs = await this.#runBatch(batch, ++this.#iterationCount); yield batchLogs; batchesLogs = batchesLogs.push(batchLogs); @@ -75,14 +76,47 @@ export class GPT extends Model<"text"> { return new EpochLogs(batchesLogs, epochTime, validation); } + async *trainNextBatches( + trainingIterator: AsyncIterator>, + maxBatchCount: number, + validationDataset?: Dataset>, + setDone?: (done: boolean) => void, + ): AsyncGenerator { + let batchesLogs = List(); + let epochTime = performance.now(); + let done = false; + + for (let batchCount = 0; batchCount < maxBatchCount; batchCount++) { + const next = await trainingIterator.next(); + if (next.done === true) { + done = true; + break; + } + + const batchLogs = await this.#runBatch(next.value, ++this.#iterationCount); + + yield batchLogs; + batchesLogs = batchesLogs.push(batchLogs); + } + + const validation = + validationDataset && (await this.evaluate(validationDataset)); + epochTime = performance.now() - epochTime; + setDone?.(done); + + return new EpochLogs(batchesLogs, epochTime, validation); + } + async #runBatch( batch: Batched, + iterationNumber: number, ): Promise { const tfBatch = this.#batchToTF(batch); let logs: tf.Logs | undefined; await this.model.fitDataset(tf.data.array([tfBatch]), { epochs: 1, + iterationOffset: iterationNumber - 1, verbose: 0, // don't pollute callbacks: { onEpochEnd: (_, cur) => { diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index a69e498b4..f11ffbbcf 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -55,9 +55,10 @@ export class GPTModel extends tf.LayersModel { : tf.train.adam(this.config.lr) } - override async fitDataset(dataset: Dataset, trainingArgs: tf.ModelFitDatasetArgs): Promise { + override async fitDataset(dataset: Dataset, trainingArgs: tf.ModelFitDatasetArgs & { iterationOffset?: number }): Promise { const callbacks = trainingArgs.callbacks as tf.CustomCallbackArgs const evalDataset = trainingArgs.validationData as tf.data.Dataset<{ xs: tf.Tensor2D, ys: tf.Tensor3D }> + const iterationOffset = trainingArgs.iterationOffset ?? 0 await callbacks.onTrainBegin?.() for (let epoch = 1; epoch <= trainingArgs.epochs; epoch++) { @@ -72,15 +73,18 @@ export class GPTModel extends tf.LayersModel { debug("after next of iterator") while (next.done !== true && iteration <= this.config.maxIter) { + const reportedIteration = iterationOffset + iteration let weightUpdateTime = performance.now() await callbacks.onEpochBegin?.(epoch) const { xs, ys } = next.value as { xs: tf.Tensor2D, ys: tf.Tensor3D } let preprocessingTime = performance.now() - debug("await batch data before {} iteration", iteration) + // debug("await batch data before {} iteration", iteration) + debug("await batch data before {} iteration", reportedIteration) // await Promise.all([xs.data(), ys.data()]) await Promise.resolve() - debug("after await batch data {} iteration", iteration) + // debug("after await batch data {} iteration", iteration) + debug("after await batch data {} iteration", reportedIteration) preprocessingTime = performance.now() - preprocessingTime // TODO include as a tensor inside the model @@ -125,7 +129,8 @@ export class GPTModel extends tf.LayersModel { if ( evalDataset !== undefined && this.config.evaluateEvery !== undefined && - iteration % this.config.evaluateEvery == 0 + // iteration % this.config.evaluateEvery == 0 + reportedIteration % this.config.evaluateEvery == 0 ){ const iterationLogs = await evaluate(this, evalDataset, this.config.maxEvalBatches) debug('evaluation metrics: %O', iterationLogs); @@ -133,7 +138,8 @@ export class GPTModel extends tf.LayersModel { const memory = tf.memory().numBytes / 1024 / 1024 / 1024 debug("training metrics: %O", { epoch, - iteration, + // iteration, + iteration: reportedIteration, loss, memory, allocated: tf.memory().numTensors, diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index fa01cdf2e..aaca2f117 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -62,6 +62,8 @@ export namespace TrainingInformation { // number of epochs between each weight sharing round. // e.g.if 3 then weights are shared every 3 epochs (in the distributed setting). roundDuration: z.number().positive().int(), + // for GPT text tasks, number of training batches between each weight sharing round. + roundIterations: z.number().positive().int().optional(), // fraction of data to keep for validation, note this only works for image data validationSplit: z.number().min(0).max(1), // batch size of training data diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 73cae1f7c..f46690b69 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -30,6 +30,15 @@ export interface RoundLogs { /** List of weight update norms */ export type WeightNormHistory = List>; +type IterationTrainableTextModel = Model<"text"> & { + trainNextBatches( + trainingIterator: AsyncIterator>, + maxBatchCount: number, + validationDataset?: Dataset>, + setDone?: (done: boolean) => void, + ): AsyncGenerator; +}; + function appendWeightHistory(weightNormHistory: WeightNormHistory, wc: number[]){ return wc.reduce((hist, t, i) => { const arr = hist.get(i, List()); @@ -53,6 +62,7 @@ export class Trainer { AsyncGenerator, RoundLogs>, void >; + readonly #roundIterations?: number; // Map of weight Index and weight update #weightNormHistory : WeightNormHistory = List(); #previousRoundWeights?: WeightsContainer; @@ -71,10 +81,18 @@ export class Trainer { this.#client = client; this.#roundDuration = task.trainingInformation.roundDuration; this.#epochs = task.trainingInformation.epochs; + this.#roundIterations = task.trainingInformation.roundIterations; if ("privacy" in task.trainingInformation) this.#privacy = task.trainingInformation.privacy; - if (!Number.isInteger(this.#epochs / this.#roundDuration)) + if (this.#roundIterations !== undefined && (task.dataType !== "text" || task.trainingInformation.tensorBackend !== "gpt")) + throw new Error("roundIterations is only supported for GPT text tasks"); + + if (this.#roundIterations !== undefined && (!Number.isInteger(this.#roundIterations) || this.#roundIterations < 1)) + throw new Error("roundIterations must be a positive integer"); + + // if (!Number.isInteger(this.#epochs / this.#roundDuration)) + if (this.#roundIterations === undefined && !Number.isInteger(this.#epochs / this.#roundDuration)) throw new Error( `round duration ${this.#roundDuration} doesn't divide number of epochs ${this.#epochs}`, ); @@ -98,7 +116,11 @@ export class Trainer { ); try { - this.#training = this.#runRounds(dataset, validationDataset); + // this.#training = this.#runRounds(dataset, validationDataset); + this.#training = + this.#roundIterations === undefined + ? this.#runRounds(dataset, validationDataset) + : this.#runIterationRounds(dataset, validationDataset); yield* this.#training; } finally { this.#training = undefined; @@ -151,6 +173,77 @@ export class Trainer { } } + async *#runIterationRounds( + dataset: Dataset>, + validationDataset?: Dataset>, + ): AsyncGenerator< + AsyncGenerator, RoundLogs>, + void + > { + if (this.#roundIterations === undefined) + throw new Error("roundIterations was not set"); + + for (let epoch = 0; epoch < this.#epochs; epoch++) { + const trainingIterator = dataset[Symbol.asyncIterator](); + let next = await trainingIterator.next(); + let pendingBatch: Batched | undefined = + next.done === true ? undefined : next.value; + + while (pendingBatch !== undefined) { + await this.#client.onRoundBeginCommunication(); + + this.#previousRoundWeights = new WeightsContainer(this.model.weights.weights.map(t => t.clone())); + + let firstBatch: Batched | undefined = pendingBatch; + pendingBatch = undefined; + let done = false; + const prefixedIterator: AsyncIterator> = { + next: async () => { + if (firstBatch !== undefined) { + const value = firstBatch; + firstBatch = undefined; + return { value, done: false }; + } + + return await trainingIterator.next(); + }, + }; + + yield this.#runIterationRound( + prefixedIterator, + this.#roundIterations, + validationDataset, + (roundDone) => done = roundDone, + ); + + let roundWeights = this.model.weights; + + if (this.#privacy !== undefined){ + const roundUpdate = roundWeights.sub(this.#previousRoundWeights); + const updateNorm = await Promise.all( + roundUpdate.weights.map(privacy.frobeniusNorm) + ); + this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); + + roundWeights = await applyOptimalPrivacy( + this.#previousRoundWeights, + roundWeights, + this.#privacy, + this.#weightNormHistory, + Number.MAX_SAFE_INTEGER, + ) + } + + const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); + this.model.weights = networkWeights; + + if (done) break; + next = await trainingIterator.next(); + pendingBatch = next.done === true ? undefined : next.value; + } + } + } + async *#runRound( dataset: Dataset>, validationDataset?: Dataset>, @@ -177,6 +270,42 @@ export class Trainer { preRoundValidation: validation, }; } + + async *#runIterationRound( + datasetIterator: AsyncIterator>, + maxBatchCount: number, + validationDataset?: Dataset>, + setDone?: (done: boolean) => void, + ): AsyncGenerator, RoundLogs> { + let epochsLogs = List(); + + debug("Run iteration-based round") + + const validation = validationDataset !== undefined ? await this.model.evaluate(validationDataset) : undefined; + + const model = this.model as unknown as IterationTrainableTextModel; + if (typeof model.trainNextBatches !== "function") + throw new Error("model does not support iteration-based training"); + + const [gen, result] = async_iterator.split( + model.trainNextBatches( + datasetIterator as AsyncIterator>, + maxBatchCount, + validationDataset as Dataset> | undefined, + setDone, + ), + ); + + yield gen; + const epochLogs = await result; + epochsLogs = epochsLogs.push(epochLogs); + + return { + epochs: epochsLogs, + participants: this.#client.nbOfParticipants, + preRoundValidation: validation, + }; + } } /** ALDP-FL implementation */ From 8d409b9cfef22d8631b51ab827f96a3a5b0f76d2 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 11 May 2026 16:53:36 +0200 Subject: [PATCH 19/89] fix mem leak --- discojs/src/models/gpt/model.ts | 1 + discojs/src/training/trainer.ts | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index f11ffbbcf..c0ec63f2b 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -117,6 +117,7 @@ export class GPTModel extends tf.LayersModel { }) const gradsClipped = clipByGlobalNormObj(grads, 1) this.optimizer.applyGradients(gradsClipped) + tf.dispose(Object.values(gradsClipped)) return lossTensor }) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index f46690b69..1a74cecbd 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -170,6 +170,9 @@ export class Trainer { // Update the local weights this.model.weights = networkWeights; + networkWeights.dispose(); + this.#previousRoundWeights.dispose(); + this.#previousRoundWeights = undefined; } } @@ -236,6 +239,9 @@ export class Trainer { const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); this.model.weights = networkWeights; + networkWeights.dispose(); + this.#previousRoundWeights.dispose(); + this.#previousRoundWeights = undefined; if (done) break; next = await trainingIterator.next(); From 144b751f0bc6c09559156214468475c1cbd1421c Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 11 May 2026 17:17:18 +0200 Subject: [PATCH 20/89] change ligs --- discojs/src/models/gpt/model.ts | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index c0ec63f2b..8be9d7379 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -79,12 +79,8 @@ export class GPTModel extends tf.LayersModel { const { xs, ys } = next.value as { xs: tf.Tensor2D, ys: tf.Tensor3D } let preprocessingTime = performance.now() - // debug("await batch data before {} iteration", iteration) - debug("await batch data before {} iteration", reportedIteration) - // await Promise.all([xs.data(), ys.data()]) - await Promise.resolve() - // debug("after await batch data {} iteration", iteration) - debug("after await batch data {} iteration", reportedIteration) + await Promise.all([xs.data(), ys.data()]) + // await Promise.resolve() preprocessingTime = performance.now() - preprocessingTime // TODO include as a tensor inside the model From c046e9edecae68ea0437e4738399e74e7d73b372 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 12 May 2026 15:50:21 +0200 Subject: [PATCH 21/89] change model to 256 --- discojs/src/default_tasks/privacyrun.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index dcdaa7640..2b743b98b 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -32,7 +32,7 @@ export const privacyrun: TaskProvider<"text", "federated"> = { batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), // contextLength: 1024, - contextLength: 512, + contextLength: 256, tensorBackend: 'gpt' } } @@ -43,7 +43,9 @@ export const privacyrun: TaskProvider<"text", "federated"> = { // The model should be in DiscoJS serialization format (created by onnx-converter) // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model.json"; - const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; + // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; + + const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_256.json"; try { const response = await fetch(modelUrl); From e89b9e0c020aab242b80298f92971d95fefe1c82 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 12 May 2026 17:47:51 +0200 Subject: [PATCH 22/89] back to 512 --- discojs/src/default_tasks/privacyrun.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index 2b743b98b..37b6b951c 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -32,7 +32,7 @@ export const privacyrun: TaskProvider<"text", "federated"> = { batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), // contextLength: 1024, - contextLength: 256, + contextLength: 512, tensorBackend: 'gpt' } } @@ -45,7 +45,7 @@ export const privacyrun: TaskProvider<"text", "federated"> = { // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; - const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_256.json"; + const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; try { const response = await fetch(modelUrl); From fbcca59330a9ad240731990dd34c333b2d0913a0 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 13 May 2026 01:31:27 +0200 Subject: [PATCH 23/89] fix end of training --- discojs/src/processing/index.spec.ts | 35 ++++++++++++++++++++++++++++ discojs/src/processing/index.ts | 1 + 2 files changed, 36 insertions(+) diff --git a/discojs/src/processing/index.spec.ts b/discojs/src/processing/index.spec.ts index 19aa2b47a..4bb3eec48 100644 --- a/discojs/src/processing/index.spec.ts +++ b/discojs/src/processing/index.spec.ts @@ -5,6 +5,12 @@ import { Dataset } from "../index.js"; import { preprocess } from "./index.js"; +async function arrayFromAsync(iter: AsyncIterable): Promise { + const ret: T[] = []; + for await (const e of iter) ret.push(e); + return ret; +} + describe("preprocess", () => { it("throws on missing column in tabular", async () => { const task: Task<"tabular", "local"> = { @@ -41,4 +47,33 @@ describe("preprocess", () => { expect(false, "should have thrown").to.be.true; }); + + it("drops incomplete text windows", async () => { + const task = { + id: "task", + dataType: "text", + displayInformation: { + title: "", + summary: { preview: "", overview: "" }, + }, + trainingInformation: { + tensorBackend: "gpt", + scheme: "local", + aggregationStrategy: "mean", + epochs: 1, + roundDuration: 1, + batchSize: 2, + validationSplit: 0, + contextLength: 4, + tokenizer: { + tokenize: () => [0, 1, 2, 3, 4, 5, 6], + }, + }, + } as unknown as Task<"text", "local">; + + const dataset = new Dataset(["ignored"]); + const preprocessed = await arrayFromAsync(preprocess(task, dataset)); + + expect(preprocessed.map(([tokens]) => tokens.size)).to.deep.equal([4]); + }); }); diff --git a/discojs/src/processing/index.ts b/discojs/src/processing/index.ts index 4011f40d1..fa73618d8 100644 --- a/discojs/src/processing/index.ts +++ b/discojs/src/processing/index.ts @@ -58,6 +58,7 @@ export function preprocess( .map((text) => tokenizer.tokenize(text)) .flatten() .batch(contextLength + 1, 1) + .filter((tokens) => tokens.size === contextLength + 1) .map((tokens) => [tokens.pop(), tokens.last()]) as Dataset< DataFormat.ModelEncoded[D] >; From dfe388ee5098dbc4a2adb264814e1964d15151b7 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 14 May 2026 19:18:11 +0200 Subject: [PATCH 24/89] fix mem script and add logs --- cli/src/measure_memorization_gpt2.ts | 119 ++++++++++++++++++++++-- discojs/src/default_tasks/privacyrun.ts | 1 + 2 files changed, 112 insertions(+), 8 deletions(-) diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts index 63f02d1ac..5f15d901d 100644 --- a/cli/src/measure_memorization_gpt2.ts +++ b/cli/src/measure_memorization_gpt2.ts @@ -14,12 +14,14 @@ interface Args { suffixLength: number; bleuThreshold: number; seed: number; + logEvery: number; savePath?: string; help?: boolean; } type PromptResult = { recordIndex: number; + recordTokenLength: number; promptLength: number; splitIndex: number; exactMatch: boolean; @@ -30,6 +32,14 @@ type PromptResult = { generatedText: string; }; +type TokenLengthStats = { + min: number; + p50: number; + p90: number; + max: number; + average: number; +}; + function parseIntegerList(raw: string): number[] { const values = raw .split(",") @@ -78,6 +88,24 @@ async function loadRecords(filePath: string, limit: number): Promise { return limit > 0 ? records.slice(0, limit) : records; } +function summarizeTokenLengths(lengths: number[]): TokenLengthStats { + if (lengths.length === 0) { + return { min: 0, p50: 0, p90: 0, max: 0, average: 0 }; + } + + const sorted = [...lengths].sort((a, b) => a - b); + const percentile = (p: number) => + sorted[Math.min(sorted.length - 1, Math.floor((sorted.length - 1) * p))]; + + return { + min: sorted[0], + p50: percentile(0.5), + p90: percentile(0.9), + max: sorted[sorted.length - 1], + average: lengths.reduce((sum, length) => sum + length, 0) / lengths.length, + }; +} + function ngrams(tokens: number[], n: number): Map { const counts = new Map(); if (tokens.length < n) return counts; @@ -197,6 +225,7 @@ async function main() { suffixLength: { type: Number, description: "Number of suffix tokens to generate and compare", defaultValue: 50 }, bleuThreshold: { type: Number, description: "BLEU threshold for approximate memorization", defaultValue: 0.75 }, seed: { type: Number, description: "Random seed for choosing record split positions", defaultValue: 42 }, + logEvery: { type: Number, description: "Print progress every N records; set 0 to disable per-record progress logs", defaultValue: 1 }, savePath: { type: String, description: "Optional JSON output path", optional: true }, help: { type: Boolean, optional: true, alias: "h", description: "Prints this usage guide" }, }, @@ -228,26 +257,85 @@ async function main() { const records = await loadRecords(args.dataPath, args.maxRecords); console.log(`Loaded ${records.length} records`); + console.log("Tokenizing records..."); + const tokenizedRecords = records.map((record) => tokenizer.tokenize(record).toArray()); + const tokenLengths = tokenizedRecords.map((ids) => ids.length); + const requiredTokensByPromptLength = Object.fromEntries( + promptLengths.map((promptLength) => [ + promptLength, + promptLength + args.suffixLength + 1, + ]), + ); + const eligibleRecordsByPromptLength = Object.fromEntries( + promptLengths.map((promptLength) => [ + promptLength, + tokenLengths.filter( + (length) => length >= promptLength + args.suffixLength + 1, + ).length, + ]), + ); + console.log("Token length stats:", summarizeTokenLengths(tokenLengths)); + console.log("Eligible records by prompt length:", eligibleRecordsByPromptLength); + console.log("Starting memorization evaluation..."); + const results: PromptResult[] = []; let skipped = 0; + const skippedByPromptLength: Record = Object.fromEntries( + promptLengths.map((promptLength) => [promptLength, 0]), + ); - for (let recordIndex = 0; recordIndex < records.length; recordIndex++) { - const record = records[recordIndex]; - const ids = tokenizer.tokenize(record).toArray(); + for (let recordIndex = 0; recordIndex < tokenizedRecords.length; recordIndex++) { + const ids = tokenizedRecords[recordIndex]; + const eligiblePromptLengths = promptLengths.filter( + (promptLength) => ids.length >= promptLength + args.suffixLength + 1, + ); + const shouldLogRecord = + args.logEvery > 0 && + (recordIndex === 0 || + (recordIndex + 1) % args.logEvery === 0 || + recordIndex === tokenizedRecords.length - 1); + + if (shouldLogRecord) { + console.log( + `Record ${recordIndex + 1}/${tokenizedRecords.length}: ${ids.length} tokens, eligible prompt lengths: ${ + eligiblePromptLengths.length > 0 ? eligiblePromptLengths.join(",") : "none" + }`, + ); + } - if (ids.length < maxPromptLength + args.suffixLength + 1) { + for (const promptLength of promptLengths) { + if (!eligiblePromptLengths.includes(promptLength)) { + skippedByPromptLength[promptLength]++; + } + } + + if (eligiblePromptLengths.length === 0) { + if (shouldLogRecord) { + console.log( + `Skipping record ${recordIndex + 1}; needs at least ${ + Math.min(...promptLengths) + args.suffixLength + 1 + } tokens for the shortest prompt/suffix setting.`, + ); + } skipped++; continue; } + const maxEligiblePromptLength = Math.max(...eligiblePromptLengths); const splitIndex = randomInt( random, - maxPromptLength, + maxEligiblePromptLength, ids.length - args.suffixLength, ); const reference = ids.slice(splitIndex, splitIndex + args.suffixLength); - for (const promptLength of promptLengths) { + for (const promptLength of eligiblePromptLengths) { + if (shouldLogRecord) { + console.log( + ` Generating ${args.suffixLength} tokens for prompt length ${promptLength} at split ${splitIndex}...`, + ); + } + const prompt = ids.slice(splitIndex - promptLength, splitIndex); const generated = await greedyGenerateGPT2( loadedModel, @@ -263,6 +351,7 @@ async function main() { results.push({ recordIndex, + recordTokenLength: ids.length, promptLength, splitIndex, exactMatch, @@ -272,10 +361,18 @@ async function main() { referenceText: tokenizer.decode(reference), generatedText: tokenizer.decode(generatedSuffix), }); + + if (shouldLogRecord) { + console.log( + ` Done prompt length ${promptLength}: exact=${exactMatch}, BLEU=${bleu.toFixed(4)}`, + ); + } } - if ((recordIndex + 1) % 10 === 0) { - console.log(`Processed ${recordIndex + 1}/${records.length} records`); + if (shouldLogRecord) { + console.log( + `Finished record ${recordIndex + 1}/${tokenizedRecords.length}; results so far: ${results.length}`, + ); } } @@ -292,9 +389,15 @@ async function main() { suffixLength: args.suffixLength, bleuThreshold: args.bleuThreshold, seed: args.seed, + logEvery: args.logEvery, modelContextLength: loadedModel.config.contextLength, }, + tokenLengthStats: summarizeTokenLengths(tokenLengths), + requiredTokensByPromptLength, + eligibleRecordsByPromptLength, skippedRecords: skipped, + skippedByPromptLength, + evaluatedRecords: new Set(results.map((result) => result.recordIndex)).size, ...summarize(results), }; diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index 37b6b951c..04b031564 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -29,6 +29,7 @@ export const privacyrun: TaskProvider<"text", "federated"> = { epochs: 1, validationSplit: 0.1, roundDuration: 1, + // Last context segment may be shorter than context length, so it will be dropped (TODO: implement padding to avoid this) batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), // contextLength: 1024, From 04a47694bff32675ef3d8c34bf2b5f2d0833be72 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 15 May 2026 15:23:46 +0200 Subject: [PATCH 25/89] debug memorization script --- cli/src/measure_memorization_gpt2.ts | 32 ++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts index 5f15d901d..6b63aa517 100644 --- a/cli/src/measure_memorization_gpt2.ts +++ b/cli/src/measure_memorization_gpt2.ts @@ -170,6 +170,8 @@ async function greedyGenerateGPT2( return output as tf.Tensor; }); + console.log("logits shape:", logits.shape); + const nextTokenTensor = tf.tidy(() => { const last = logits.slice([0, modelInput.length - 1, 0], [1, 1, -1]); return last.squeeze().argMax(); @@ -297,8 +299,7 @@ async function main() { if (shouldLogRecord) { console.log( - `Record ${recordIndex + 1}/${tokenizedRecords.length}: ${ids.length} tokens, eligible prompt lengths: ${ - eligiblePromptLengths.length > 0 ? eligiblePromptLengths.join(",") : "none" + `Record ${recordIndex + 1}/${tokenizedRecords.length}: ${ids.length} tokens, eligible prompt lengths: ${eligiblePromptLengths.length > 0 ? eligiblePromptLengths.join(",") : "none" }`, ); } @@ -312,8 +313,7 @@ async function main() { if (eligiblePromptLengths.length === 0) { if (shouldLogRecord) { console.log( - `Skipping record ${recordIndex + 1}; needs at least ${ - Math.min(...promptLengths) + args.suffixLength + 1 + `Skipping record ${recordIndex + 1}; needs at least ${Math.min(...promptLengths) + args.suffixLength + 1 } tokens for the shortest prompt/suffix setting.`, ); } @@ -344,6 +344,30 @@ async function main() { loadedModel.config.contextLength, ); const generatedSuffix = generated.slice(prompt.length, prompt.length + args.suffixLength); + + console.log("================================"); + console.log("PROMPT LENGTH:", promptLength); + + console.log("\nPROMPT IDS:"); + console.log(prompt.slice(0, 30)); + + console.log("\nGENERATED IDS:"); + console.log(generatedSuffix.slice(0, 30)); + + console.log("\nREFERENCE IDS:"); + console.log(reference.slice(0, 30)); + + console.log("\nPROMPT TEXT:"); + console.log(JSON.stringify(tokenizer.decode(prompt))); + + console.log("\nGENERATED TEXT:"); + console.log(JSON.stringify(tokenizer.decode(generatedSuffix))); + + console.log("\nREFERENCE TEXT:"); + console.log(JSON.stringify(tokenizer.decode(reference))); + + console.log("================================"); + const exactMatch = generatedSuffix.length === reference.length && generatedSuffix.every((token, i) => token === reference[i]); From e2d23610b33d35f2680bb9fb951a116f9e7ab59c Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 16 May 2026 15:42:00 +0200 Subject: [PATCH 26/89] change benchmark --- cli/src/evaluate_finetuned_gpt2.ts | 81 +++++++++++++----------------- 1 file changed, 36 insertions(+), 45 deletions(-) diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index ecccc6bc1..7aff1421c 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -55,68 +55,47 @@ function parseSample(sample: string) { for (const line of lines) { if (line.startsWith("Answer:")) { - answer = line.replace("Answer:", "").trim(); + // answer = line.replace("Answer:", "").trim(); + answer = line.replace("Answer:", "").trim().charAt(0).toUpperCase(); } else { promptLines.push(line); } } - const basePrompt = promptLines.join("\n"); + const basePrompt = promptLines.join("\n").trim(); return { basePrompt, answer }; } -// ========================= -// SOFTMAX (for safety) -// ========================= -async function scoreText( +async function scoreContinuation( tfModel: tf.LayersModel, tokenizer: Tokenizer, - text: string + prompt: string, + continuation: string ): Promise { - const tokens = tokenizer.tokenize(text); - - if (tokens.size < 2) return -Infinity; + const promptTokens = tokenizer.tokenize(prompt).toArray(); + const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); - const inputTokens = tokens.slice(0, tokens.size - 1).toArray(); - const targets = tokens.slice(1).toArray(); - - const inputTensor = tf.tensor([inputTokens], [1, inputTokens.length], "int32"); + const inputTokens = fullTokens.slice(0, -1); + const inputTensor = tf.tensor2d([inputTokens], [1, inputTokens.length], "int32"); const logits = tfModel.predict(inputTensor) as tf.Tensor; - const logitsArray = await logits.array() as number[][][]; + const logProbs = tf.logSoftmax(logits, -1); + const arr = await logProbs.array() as number[][][]; let score = 0; + let count = 0; - for (let i = 0; i < targets.length; i++) { - const stepLogits = logitsArray[0][i]; - - const logit = stepLogits[targets[i]] ?? -100; - - score += logit; + for (let targetPos = promptTokens.length; targetPos < fullTokens.length; targetPos++) { + const targetToken = fullTokens[targetPos]; + score += arr[0][targetPos - 1][targetToken]; + count++; } inputTensor.dispose(); logits.dispose(); + logProbs.dispose(); - return score; -} - -// ========================= -// SCORE OPTIONS -// ========================= -async function scoreOptions( - tfModel: tf.LayersModel, - tokenizer: Tokenizer, - texts: string[] -): Promise { - const scores: number[] = []; - - for (const t of texts) { - const s = await scoreText(tfModel, tokenizer, t); - scores.push(s); - } - - return scores; + return score / count; } // ========================= @@ -148,6 +127,7 @@ async function benchmarkQA( predicted: string; answer: string; correct: boolean; + scores: Record; }; const logs: PredictionLog[] = []; @@ -157,11 +137,17 @@ async function benchmarkQA( for (const sample of dataset) { const { basePrompt, answer } = parseSample(sample); - const texts = options.map( - (opt) => `${basePrompt}\nAnswer: ${opt}` - ); + if (!options.includes(answer)) { + console.log("Invalid answer:", JSON.stringify(answer)); + continue; + } + + const prompt = `${basePrompt}\nAnswer:`; - const scores = await scoreOptions(tfModel, tokenizer, texts); + const scores: number[] = []; + for (const opt of options) { + scores.push(await scoreContinuation(tfModel, tokenizer, prompt, ` ${opt}`)); + } let bestIdx = 0; for (let i = 1; i < scores.length; i++) { @@ -177,10 +163,15 @@ async function benchmarkQA( confusion[answer][predicted]++; } + const scoreMap = Object.fromEntries( + options.map((opt, i) => [opt, scores[i]]) + ); + logs.push({ predicted, answer, - correct: predicted === answer + correct: predicted === answer, + scores: scoreMap }); if (total % 50 === 0) { From 25fcf2822bb13096efa4edefc12fdfe9ac738d71 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 23 May 2026 18:36:20 +0200 Subject: [PATCH 27/89] add client id on debug logs, and final aggregation wait between clients --- cli/src/args.ts | 11 ++++++++++- cli/src/cli.ts | 6 +++++- discojs/src/models/gpt/index.ts | 4 ++++ discojs/src/models/gpt/model.ts | 19 ++++++++++++++----- discojs/src/task/training_information.ts | 2 ++ discojs/src/training/disco.ts | 16 +++++++++++++++- discojs/src/training/trainer.ts | 17 +++++++++++++++-- 7 files changed, 65 insertions(+), 10 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index f7755b43d..f529011bb 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -22,6 +22,7 @@ export interface BenchmarkArguments { roundIterations?: number batchSize: number validationSplit: number + validationFrequency?: number datasetPath?: string validationDatasetPath?: string outputPath?: string @@ -63,6 +64,7 @@ const unsafeArgs = parse( roundIterations: { type: Number, description: 'For GPT text tasks, aggregate every N training batches without rewinding the dataset', optional: true }, batchSize: { type: Number, alias: 'b', description: 'Training batch size', defaultValue: 10 }, validationSplit : { type: Number, alias: 'v', description: 'Validation dataset ratio', defaultValue: 0.2 }, + 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 ./', optional: true }, @@ -135,7 +137,14 @@ export const args: BenchmarkArguments = { task.trainingInformation.roundDuration = unsafeArgs.roundDuration; task.trainingInformation.epochs = unsafeArgs.epochs; task.trainingInformation.validationSplit = unsafeArgs.validationSplit; - (task.trainingInformation as typeof task.trainingInformation & { roundIterations?: number }).roundIterations = unsafeArgs.roundIterations; + (task.trainingInformation as typeof task.trainingInformation & { + roundIterations?: number; + validationFrequency?: number; + }).roundIterations = unsafeArgs.roundIterations; + (task.trainingInformation as typeof task.trainingInformation & { + roundIterations?: number; + validationFrequency?: number; + }).validationFrequency = unsafeArgs.validationFrequency; const {aggregator, clippingRadius, maxIterations, beta, maxShareValue} = unsafeArgs; diff --git a/cli/src/cli.ts b/cli/src/cli.ts index e780f6b65..9cd89a8d6 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -43,7 +43,11 @@ async function runUser( const trainingScheme = task.trainingInformation.scheme as N const aggregator = aggregators.getAggregator(task) const client = clients.getClient(trainingScheme, url, task, aggregator) - const disco = new Disco(task, client, { scheme: trainingScheme, preprocessOnce: false }); + const disco = new Disco(task, client, { + scheme: trainingScheme, + preprocessOnce: false, + debugLabel: `client${userIndex}`, + }); // For local training, load model from provider before training starts // if (trainingScheme === "local") { diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 458dad7e6..b214a644b 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -44,6 +44,10 @@ export class GPT extends Model<"text"> { this.#vocabSize = partialConfig?.vocabSize ?? DefaultGPTConfig.vocabSize; } + setDebugLabel(label: string): void { + this.model.setDebugLabel(label); + } + /** * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs) * This allows for getting logs and stopping training without callbacks. diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 8be9d7379..66f399a3f 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -27,6 +27,7 @@ export declare abstract class Dataset { */ export class GPTModel extends tf.LayersModel { protected readonly config: Required + #debugLabel?: string constructor(partialConfig?: Partial, layersModel?: tf.LayersModel) { // Fill missing config parameters with default values @@ -48,6 +49,14 @@ export class GPTModel extends tf.LayersModel { return this.config } + setDebugLabel(label: string): void { + this.#debugLabel = label + } + + #debugMessage(message: string): string { + return this.#debugLabel === undefined ? message : `[${this.#debugLabel}] ${message}` + } + override compile() { if (this.optimizer !== undefined) return this.optimizer = this.config.weightDecay !== 0 @@ -66,11 +75,11 @@ export class GPTModel extends tf.LayersModel { let averageLoss = 0 let iteration = 1 - debug("before iterator init") + debug(this.#debugMessage("before iterator init")) const iterator = await dataset.iterator() - debug("after getting iterator, before next") + debug(this.#debugMessage("after getting iterator, before next")) let next = await iterator.next() - debug("after next of iterator") + debug(this.#debugMessage("after next of iterator")) while (next.done !== true && iteration <= this.config.maxIter) { const reportedIteration = iterationOffset + iteration @@ -130,10 +139,10 @@ export class GPTModel extends tf.LayersModel { reportedIteration % this.config.evaluateEvery == 0 ){ const iterationLogs = await evaluate(this, evalDataset, this.config.maxEvalBatches) - debug('evaluation metrics: %O', iterationLogs); + debug(this.#debugMessage('evaluation metrics: %O'), iterationLogs); } const memory = tf.memory().numBytes / 1024 / 1024 / 1024 - debug("training metrics: %O", { + debug(this.#debugMessage("training metrics: %O"), { epoch, // iteration, iteration: reportedIteration, diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index aaca2f117..a1f7706dd 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -64,6 +64,8 @@ export namespace TrainingInformation { roundDuration: z.number().positive().int(), // for GPT text tasks, number of training batches between each weight sharing round. roundIterations: z.number().positive().int().optional(), + // run validation every N aggregation rounds. If 0, validation metrics are skipped. + validationFrequency: z.number().nonnegative().int().optional(), // fraction of data to keep for validation, note this only works for image data validationSplit: z.number().min(0).max(1), // batch size of training data diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 33c63c244..e1b7f3e70 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -30,6 +30,7 @@ const debug = createDebug("discojs:training:disco"); interface DiscoConfig { scheme: N; logger: Logger; + debugLabel?: string; /** * keep preprocessed dataset in memory while training @@ -88,6 +89,7 @@ export class Disco extends EventEmitter<{ readonly #logger: Logger; readonly #task: Task; readonly #preprocessOnce: boolean; + readonly #debugLabel?: string; /** * Connect to the given task and get ready to train. @@ -102,7 +104,7 @@ export class Disco extends EventEmitter<{ config: Partial>, ) { super(); - const { scheme, logger, preprocessOnce } = { + const { scheme, logger, preprocessOnce, debugLabel } = { // cast as typescript is bad at generic scheme: task.trainingInformation.scheme as N, logger: new ConsoleLogger(), @@ -128,6 +130,7 @@ export class Disco extends EventEmitter<{ this.#logger = logger; this.#preprocessOnce = preprocessOnce; + this.#debugLabel = debugLabel; this.#client = client; this.#task = task; this.trainer = new Trainer(task, client); @@ -226,6 +229,7 @@ export class Disco extends EventEmitter<{ // TODO unsafe cast debug("Connecting to client and fetching initial model..."); this.trainer.model = (await this.#client.connect()) as Model; + this.#setModelDebugLabel(this.trainer.model); debug("Initial model fetched successfully"); if (this.trainer.model === null) { debug(`No pre-trained model provided for client, initializing randomly...`); @@ -287,6 +291,16 @@ export class Disco extends EventEmitter<{ await this.#client.disconnect(); } + #setModelDebugLabel(model: Model): void { + if (this.#debugLabel === undefined) return; + + const labeledModel = model as Model & { + setDebugLabel?: (label: string) => void; + }; + + labeledModel.setDebugLabel?.(this.#debugLabel); + } + async #preprocessSplitAndBatch( dataset: Dataset, ): Promise< diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 1a74cecbd..5687eac2d 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -63,6 +63,7 @@ export class Trainer { void >; readonly #roundIterations?: number; + readonly #validationFrequency?: number; // Map of weight Index and weight update #weightNormHistory : WeightNormHistory = List(); #previousRoundWeights?: WeightsContainer; @@ -82,6 +83,7 @@ export class Trainer { this.#roundDuration = task.trainingInformation.roundDuration; this.#epochs = task.trainingInformation.epochs; this.#roundIterations = task.trainingInformation.roundIterations; + this.#validationFrequency = task.trainingInformation.validationFrequency; if ("privacy" in task.trainingInformation) this.#privacy = task.trainingInformation.privacy; @@ -91,6 +93,9 @@ export class Trainer { if (this.#roundIterations !== undefined && (!Number.isInteger(this.#roundIterations) || this.#roundIterations < 1)) throw new Error("roundIterations must be a positive integer"); + if (this.#validationFrequency !== undefined && (!Number.isInteger(this.#validationFrequency) || this.#validationFrequency < 0)) + throw new Error("validationFrequency must be a non-negative integer"); + // if (!Number.isInteger(this.#epochs / this.#roundDuration)) if (this.#roundIterations === undefined && !Number.isInteger(this.#epochs / this.#roundDuration)) throw new Error( @@ -145,7 +150,7 @@ export class Trainer { // Store the clean weight before starting the communication this.#previousRoundWeights = new WeightsContainer(this.model.weights.weights.map(t => t.clone())); - yield this.#runRound(dataset, validationDataset); + yield this.#runRound(dataset, this.#shouldValidateRound(round) ? validationDataset : undefined); let roundWeights = this.model.weights; @@ -186,6 +191,7 @@ export class Trainer { if (this.#roundIterations === undefined) throw new Error("roundIterations was not set"); + let round = 0; for (let epoch = 0; epoch < this.#epochs; epoch++) { const trainingIterator = dataset[Symbol.asyncIterator](); let next = await trainingIterator.next(); @@ -215,7 +221,7 @@ export class Trainer { yield this.#runIterationRound( prefixedIterator, this.#roundIterations, - validationDataset, + this.#shouldValidateRound(round) ? validationDataset : undefined, (roundDone) => done = roundDone, ); @@ -243,6 +249,7 @@ export class Trainer { this.#previousRoundWeights.dispose(); this.#previousRoundWeights = undefined; + round++; if (done) break; next = await trainingIterator.next(); pendingBatch = next.done === true ? undefined : next.value; @@ -312,6 +319,12 @@ export class Trainer { preRoundValidation: validation, }; } + + #shouldValidateRound(round: number): boolean { + if (this.#validationFrequency === undefined) return true; + if (this.#validationFrequency === 0) return false; + return round % this.#validationFrequency === 0; + } } /** ALDP-FL implementation */ From 3f6755ca002db707a3f0ce291dc9ba3b5af941a2 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 28 May 2026 13:24:11 +0200 Subject: [PATCH 28/89] Implement goldfish loss and make benchmark and memorizarion script junk free --- cli/src/args.ts | 24 +++++++ cli/src/evaluate_finetuned_gpt2.ts | 15 ----- discojs/src/index.ts | 1 + discojs/src/models/gpt/config.ts | 7 ++ discojs/src/models/gpt/index.ts | 6 +- discojs/src/models/gpt/model.ts | 82 +++++++++++++++++++++++- discojs/src/models/index.ts | 1 + discojs/src/task/training_information.ts | 9 +++ discojs/src/training/disco.ts | 12 ++++ 9 files changed, 138 insertions(+), 19 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index f529011bb..77a59e485 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -26,6 +26,10 @@ export interface BenchmarkArguments { datasetPath?: string validationDatasetPath?: string outputPath?: string + goldfishLoss: boolean + goldfishK: number + goldfishH: number + goldfishPadTokenId?: number // DP epsilon?: number @@ -68,6 +72,10 @@ const unsafeArgs = parse( 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 ./', 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 }, 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 }, host: { @@ -146,6 +154,22 @@ export const args: BenchmarkArguments = { validationFrequency?: number; }).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, + }; + } + const {aggregator, clippingRadius, maxIterations, beta, maxShareValue} = unsafeArgs; // For aggregators diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 7aff1421c..84f540b11 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -13,14 +13,9 @@ interface Args { help?: boolean; } -// ========================= // HOW TO RUN -// ========================= // npm -w cli run eval_finetuned_gpt2 -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/train_no_exp.txt --maxSamples 100 -// ========================= -// LOAD DATASET -// ========================= async function loadDataset(filePath: string, limit = -1): Promise { const text = await fs.readFile(filePath, "utf-8"); const lines = text.split("\n"); @@ -44,9 +39,6 @@ async function loadDataset(filePath: string, limit = -1): Promise { return samples; } -// ========================= -// PARSE SAMPLE -// ========================= function parseSample(sample: string) { const lines = sample.split("\n"); @@ -55,7 +47,6 @@ function parseSample(sample: string) { for (const line of lines) { if (line.startsWith("Answer:")) { - // answer = line.replace("Answer:", "").trim(); answer = line.replace("Answer:", "").trim().charAt(0).toUpperCase(); } else { promptLines.push(line); @@ -98,9 +89,6 @@ async function scoreContinuation( return score / count; } -// ========================= -// BENCHMARK -// ========================= async function benchmarkQA( model: models.GPT, tokenizer: Tokenizer, @@ -205,9 +193,6 @@ async function benchmarkQA( } } -// ========================= -// MAIN -// ========================= async function main() { const args = parse({ modelPath: { type: String }, diff --git a/discojs/src/index.ts b/discojs/src/index.ts index 188090b0e..1ccd6a9d5 100644 --- a/discojs/src/index.ts +++ b/discojs/src/index.ts @@ -15,6 +15,7 @@ export { Model, BatchLogs, EpochLogs, + GoldfishLossConfig, Tokenizer, ValidationMetrics, } from "./models/index.js"; diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index b2a86340f..c4e68bbe6 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -27,6 +27,13 @@ export type GPTConfig = { nEmbd?: number seed?: number, } + +export type GoldfishLossConfig = { + enabled: boolean + k: number + h: number + padTokenId?: number +} // for a benchmark of performance, see https://github.com/epfml/disco/pull/659 export const DefaultGPTConfig: Required = { lr: 0.001, diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index b214a644b..3d480d2c0 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -15,7 +15,7 @@ import { BatchLogs, Model, EpochLogs } from "../index.js"; import { GPTModel } from "./model.js"; import evaluate from "./evaluate.js"; import { DefaultGPTConfig, DefaultGenerationConfig } from './config.js' -import type { GPTConfig, GenerationConfig } from './config.js' +import type { GoldfishLossConfig, GPTConfig, GenerationConfig } from './config.js' const debug = createDebug("discojs:models:gpt"); @@ -48,6 +48,10 @@ export class GPT extends Model<"text"> { this.model.setDebugLabel(label); } + setGoldfishLoss(config: GoldfishLossConfig | undefined): void { + this.model.setGoldfishLoss(config); + } + /** * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs) * This allows for getting logs and stopping training without callbacks. diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 66f399a3f..cb8571b21 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -1,7 +1,7 @@ import createDebug from "debug"; import * as tf from '@tensorflow/tfjs' -import type { GPTConfig } from './config.js' +import type { GoldfishLossConfig, GPTConfig } from './config.js' import { getModelSizes, DefaultGPTConfig } from './config.js' import { getCustomAdam, clipByGlobalNormObj } from './optimizers.js' import evaluate from './evaluate.js' @@ -28,6 +28,7 @@ export declare abstract class Dataset { export class GPTModel extends tf.LayersModel { protected readonly config: Required #debugLabel?: string + #goldfishLoss?: GoldfishLossConfig constructor(partialConfig?: Partial, layersModel?: tf.LayersModel) { // Fill missing config parameters with default values @@ -53,6 +54,10 @@ export class GPTModel extends tf.LayersModel { this.#debugLabel = label } + setGoldfishLoss(config: GoldfishLossConfig | undefined): void { + this.#goldfishLoss = config?.enabled === true ? config : undefined + } + #debugMessage(message: string): string { return this.#debugLabel === undefined ? message : `[${this.#debugLabel}] ${message}` } @@ -111,6 +116,12 @@ export class GPTModel extends tf.LayersModel { // tf.dispose([accTensor]) accuracyFraction = [Number.NaN, Number.NaN]; + const goldfishLoss = this.#goldfishLoss + const goldfishMask = + goldfishLoss === undefined + ? undefined + : this.#buildGoldfishMask(xs, goldfishLoss) + const lossTensor = tf.tidy(() => { const { grads, value: lossTensor } = this.optimizer.computeGradients(() => { const logits = this.apply(xs) @@ -118,13 +129,16 @@ export class GPTModel extends tf.LayersModel { throw new Error('model outputs too many tensor') if (logits instanceof tf.SymbolicTensor) throw new Error('model outputs symbolic tensor') - return tf.losses.softmaxCrossEntropy(ys, logits) + return goldfishMask === undefined || goldfishLoss === undefined + ? tf.losses.softmaxCrossEntropy(ys, logits) + : this.#goldfishLossTensor(ys, logits, goldfishMask, goldfishLoss) }) const gradsClipped = clipByGlobalNormObj(grads, 1) this.optimizer.applyGradients(gradsClipped) tf.dispose(Object.values(gradsClipped)) return lossTensor }) + goldfishMask?.dispose() const loss = await lossTensor.array() averageLoss += loss @@ -144,7 +158,6 @@ export class GPTModel extends tf.LayersModel { const memory = tf.memory().numBytes / 1024 / 1024 / 1024 debug(this.#debugMessage("training metrics: %O"), { epoch, - // iteration, iteration: reportedIteration, loss, memory, @@ -172,4 +185,67 @@ export class GPTModel extends tf.LayersModel { await callbacks.onTrainEnd?.() return new tf.History() } + + #goldfishLossTensor( + ys: tf.Tensor3D, + logits: tf.Tensor | tf.Tensor[], + goldfishMask: tf.Tensor2D, + config: GoldfishLossConfig, + ): tf.Scalar { + if (Array.isArray(logits)) + throw new Error('model outputs too many tensor') + if (logits.rank !== 3) + throw new Error('model outputs wrong shape') + + const tokenLosses = tf.neg( + tf.sum(tf.mul(ys, tf.logSoftmax(logits as tf.Tensor3D, -1)), -1), + ) as tf.Tensor2D + + const supervisedMask = + config.padTokenId === undefined + ? goldfishMask + : tf.mul( + goldfishMask, + tf.cast(tf.notEqual(tf.argMax(ys, -1), config.padTokenId), 'float32'), + ) + + const denominator = tf.maximum(tf.sum(supervisedMask), tf.scalar(1)) + return tf.div(tf.sum(tf.mul(tokenLosses, supervisedMask)), denominator) + } + + #buildGoldfishMask( + inputIds: tf.Tensor2D, + config: GoldfishLossConfig, + ): tf.Tensor2D { + const rows = inputIds.arraySync() + const mask = rows.map((row) => + row.map((_, targetOffset) => { + const targetIndex = targetOffset + 1 + const start = Math.max(0, targetIndex - config.h) + const context = row.slice(start, targetIndex) + return this.#hashTokenContext(context) % config.k === 0 ? 0 : 1 + }), + ) + + return tf.tensor2d(mask, inputIds.shape, 'float32') + } + + #hashTokenContext(tokens: number[]): number { + let hash = 0x811c9dc5 + + for (const token of tokens) { + hash ^= token & 0xff + hash = Math.imul(hash, 0x01000193) + hash ^= (token >>> 8) & 0xff + hash = Math.imul(hash, 0x01000193) + hash ^= (token >>> 16) & 0xff + hash = Math.imul(hash, 0x01000193) + hash ^= (token >>> 24) & 0xff + hash = Math.imul(hash, 0x01000193) + hash ^= 0xff + hash = Math.imul(hash, 0x01000193) + } + + return hash >>> 0 + } } diff --git a/discojs/src/models/index.ts b/discojs/src/models/index.ts index 96d253f0b..4db0cdf0a 100644 --- a/discojs/src/models/index.ts +++ b/discojs/src/models/index.ts @@ -5,6 +5,7 @@ export { Tokenizer } from "./tokenizer.js"; export { GPT } from './gpt/index.js' export { ONNXModel } from './onnx.js' export { GPTConfig } from './gpt/config.js' +export { GoldfishLossConfig } from './gpt/config.js' export { evaluate as evaluate_hellaswag, HellaSwagDataset, diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index a1f7706dd..c9c74dc6d 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -96,6 +96,15 @@ export namespace TrainingInformation { // the maximum length of a input string used as input to a GPT model. It is used during preprocessing to // truncate strings to a maximum length. The default value is tokenizer.model_max_length contextLength: z.number().positive().int(), + // Goldfish loss drops a deterministic subset of shifted target-token losses while keeping full inputs. + goldfishLoss: z + .object({ + enabled: z.boolean(), + k: z.number().positive().int().default(4), + h: z.number().positive().int().default(13), + padTokenId: z.number().int().optional(), + }) + .optional(), }), } satisfies Record; diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index e1b7f3e70..0eca48cf3 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -15,6 +15,7 @@ import type { Model, Network, Task, + GoldfishLossConfig, } from "../index.js"; import type { Aggregator } from "../aggregator/index.js"; import { getAggregator } from "../aggregator/index.js"; @@ -230,6 +231,7 @@ export class Disco extends EventEmitter<{ debug("Connecting to client and fetching initial model..."); this.trainer.model = (await this.#client.connect()) as Model; this.#setModelDebugLabel(this.trainer.model); + this.#setModelTrainingOptions(this.trainer.model); debug("Initial model fetched successfully"); if (this.trainer.model === null) { debug(`No pre-trained model provided for client, initializing randomly...`); @@ -301,6 +303,16 @@ export class Disco extends EventEmitter<{ labeledModel.setDebugLabel?.(this.#debugLabel); } + #setModelTrainingOptions(model: Model): void { + if (this.#task.dataType !== "text") return; + + const configurableModel = model as Model & { + setGoldfishLoss?: (config: GoldfishLossConfig | undefined) => void; + }; + + configurableModel.setGoldfishLoss?.(this.#task.trainingInformation.goldfishLoss); + } + async #preprocessSplitAndBatch( dataset: Dataset, ): Promise< From d15da843e13c6c7e938724c607a5df2a840c955d Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 1 Jun 2026 17:49:03 +0200 Subject: [PATCH 29/89] add def task for local finetuning --- cli/src/args.ts | 1 + .../centralized_gpt2_finetune.ts | 69 +++++++++++++++++++ discojs/src/default_tasks/index.ts | 3 +- 3 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 discojs/src/default_tasks/centralized_gpt2_finetune.ts diff --git a/cli/src/args.ts b/cli/src/args.ts index 77a59e485..ec492fbfd 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -119,6 +119,7 @@ const supportedTasks = Map( defaultTasks.tinderDog, defaultTasks.mnist, defaultTasks.privacyrun, + defaultTasks.centralizedGPT2FineTune, ).map( async (t) => [(await t.getTask()).id, t] as [ diff --git a/discojs/src/default_tasks/centralized_gpt2_finetune.ts b/discojs/src/default_tasks/centralized_gpt2_finetune.ts new file mode 100644 index 000000000..833a1622a --- /dev/null +++ b/discojs/src/default_tasks/centralized_gpt2_finetune.ts @@ -0,0 +1,69 @@ +import type { TaskProvider } from "../index.js"; +import { Tokenizer, models, serialization } from "../index.js"; + +export const centralizedGPT2FineTune: TaskProvider<"text", "local"> = { + async getTask() { + return { + id: 'centralized-gpt2-finetune', + dataType: "text", + displayInformation: { + title: "GPT Privacy-Preserving Fine-tuning", + summary: { + preview: 'Fine-tune a pre-trained GPT model collaboratively and privately.', + overview: "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning." + }, + model: [ + "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", + "The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer.", + "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", + "Context length is kept at 1024 to match the pre-trained model, with batch size at 1.", + ].join(" "), + dataFormatInformation: 'You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.', + dataExample: + "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", + }, + trainingInformation: { + scheme: 'local', + aggregationStrategy: 'mean', + epochs: 1, + validationSplit: 0.1, + roundDuration: 1, + // Last context segment may be shorter than context length, so it will be dropped (TODO: implement padding to avoid this) + batchSize: 8, + tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), + contextLength: 512, + tensorBackend: 'gpt' + } + } + }, + + async getModel() { + // Load the pre-trained ONNX-converted model from Google Cloud Storage + // The model should be in DiscoJS serialization format (created by onnx-converter) + + const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; + + try { + const response = await fetch(modelUrl); + if (!response.ok) { + throw new Error(`HTTP ${response.status}: ${response.statusText}`); + } + + const arrayBuffer = await response.arrayBuffer(); + const encodedData = new Uint8Array(arrayBuffer); + + const model = await serialization.model.decode(encodedData); + + if (!(model instanceof models.GPT)) { + throw new Error("Loaded model is not a GPT model"); + } + + console.log("Successfully loaded pre-trained GPT model from Google Cloud Storage"); + + return model; + } catch (error) { + console.error("Failed to load model from Google Cloud Storage:", error); + throw new Error(`Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`); + } + }, +} diff --git a/discojs/src/default_tasks/index.ts b/discojs/src/default_tasks/index.ts index f46f27026..94c536412 100644 --- a/discojs/src/default_tasks/index.ts +++ b/discojs/src/default_tasks/index.ts @@ -5,4 +5,5 @@ export { simpleFace } from './simple_face.js' export { titanic } from './titanic.js' export { wikitext } from './wikitext.js' export { tinderDog } from './tinder_dog.js' -export { privacyrun } from './privacyrun.js' \ No newline at end of file +export { privacyrun } from './privacyrun.js' +export { centralizedGPT2FineTune } from './centralized_gpt2_finetune.js' \ No newline at end of file From 05cf2390ccd88908d881a706b58fc615197cefce Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 1 Jun 2026 18:53:21 +0200 Subject: [PATCH 30/89] fix bug --- cli/src/data.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cli/src/data.ts b/cli/src/data.ts index a1f41e9aa..c5c482ba3 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -171,7 +171,8 @@ export async function getTaskData( case "mnist_federated": case "mnist": return loadData("mnist", userIdx) as Dataset; - case "privacyrun": { + case "privacyrun": + case "centralized-gpt2-finetune": { const filePath = isValidation && validationDatasetPath ? validationDatasetPath From c0c41a1dfa88baf597cd683ce32035258fbf9cec Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 1 Jun 2026 19:24:20 +0200 Subject: [PATCH 31/89] bug fix --- discojs/src/client/local_client.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/discojs/src/client/local_client.ts b/discojs/src/client/local_client.ts index 4ef477f2a..9a204e3eb 100644 --- a/discojs/src/client/local_client.ts +++ b/discojs/src/client/local_client.ts @@ -10,8 +10,9 @@ export class LocalClient extends Client<"local"> { override onRoundBeginCommunication(): Promise { return Promise.resolve(); } - // Simply return the local weights + // Return clones so the trainer can dispose the communication result without + // disposing tensors owned by the model. override onRoundEndCommunication(weights: WeightsContainer): Promise { - return Promise.resolve(weights); + return Promise.resolve(weights.map((weight) => weight.clone())); } } From 95197f72abc234a06019dacffe5dffd8beaf9b50 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 1 Jun 2026 20:22:00 +0200 Subject: [PATCH 32/89] add validation before and after aggreagtion --- discojs/src/training/disco.ts | 12 +++- discojs/src/training/trainer.ts | 121 ++++++++++++++++++-------------- 2 files changed, 79 insertions(+), 54 deletions(-) diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 0eca48cf3..2e2629883 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -53,7 +53,9 @@ export type SummaryLogs = { roundValidationLoss?: number, roundValidationAccuracy?: number, validationLoss?: number, - validationAccuracy?: number + validationAccuracy?: number, + postAggregationValidationLoss?: number, + postAggregationValidationAccuracy?: number } export type RoundStatus = 'not enough participants' | // Server notification to wait for more participants @@ -73,6 +75,8 @@ function buildSummaryLog(roundNum: number, epochNum: number, roundLogs: RoundLog roundValidationAccuracy: roundLogs.preRoundValidation?.accuracy, validationLoss: epochLogs.validation?.loss, validationAccuracy: epochLogs.validation?.accuracy, + postAggregationValidationLoss: roundLogs.postAggregationValidation?.loss, + postAggregationValidationAccuracy: roundLogs.postAggregationValidation?.accuracy, } } @@ -259,6 +263,8 @@ export class Disco extends EventEmitter<{ `Round: ${roundNum}`, `Initial round loss: ${roundLogs.preRoundValidation?.loss}`, `Initial round accuracy: ${roundLogs.preRoundValidation?.accuracy}`, + `Post-aggregation loss: ${roundLogs.postAggregationValidation?.loss}`, + `Post-aggregation accuracy: ${roundLogs.postAggregationValidation?.accuracy}`, ].join("\n"), ); @@ -271,10 +277,10 @@ export class Disco extends EventEmitter<{ ` Training accuracy: ${epochLogs.training.accuracy}`, ` Peak memory: ${epochLogs.peakMemory}`, epochLogs.validation !== undefined - ? ` Validation loss: ${epochLogs.validation.loss}` + ? ` Pre-aggregation validation loss: ${epochLogs.validation.loss}` : "", epochLogs.validation !== undefined - ? ` Validation accuracy: ${epochLogs.validation.accuracy}` + ? ` Pre-aggregation validation accuracy: ${epochLogs.validation.accuracy}` : "", ].join("\n"), ); diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 5687eac2d..ed0a72c7e 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -25,6 +25,7 @@ export interface RoundLogs { epochs: List; participants: number; preRoundValidation?: ValidationMetrics; + postAggregationValidation?: ValidationMetrics; } /** List of weight update norms */ @@ -67,6 +68,7 @@ export class Trainer { // Map of weight Index and weight update #weightNormHistory : WeightNormHistory = List(); #previousRoundWeights?: WeightsContainer; + readonly #shouldValidateAfterAggregation: boolean; public get model(): Model { if (this.#model === undefined) @@ -84,6 +86,7 @@ export class Trainer { this.#epochs = task.trainingInformation.epochs; this.#roundIterations = task.trainingInformation.roundIterations; this.#validationFrequency = task.trainingInformation.validationFrequency; + this.#shouldValidateAfterAggregation = task.trainingInformation.scheme !== "local"; if ("privacy" in task.trainingInformation) this.#privacy = task.trainingInformation.privacy; @@ -150,34 +153,18 @@ export class Trainer { // Store the clean weight before starting the communication this.#previousRoundWeights = new WeightsContainer(this.model.weights.weights.map(t => t.clone())); - yield this.#runRound(dataset, this.#shouldValidateRound(round) ? validationDataset : undefined); + const roundValidationDataset = this.#shouldValidateRound(round) + ? validationDataset + : undefined; - let roundWeights = this.model.weights; - - // Apply differential privacy before sharing the weight updates with other nodes - if (this.#privacy !== undefined){ - const roundUpdate = roundWeights.sub(this.#previousRoundWeights); - const updateNorm = await Promise.all( - roundUpdate.weights.map(privacy.frobeniusNorm) - ); - this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); - - roundWeights = await applyOptimalPrivacy( - this.#previousRoundWeights, - roundWeights, - this.#privacy, - this.#weightNormHistory, + yield this.#runRound( + dataset, + roundValidationDataset, + async () => this.#finishRoundCommunication( totalRound, - ) - } - // Get the updated weights - const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); - - // Update the local weights - this.model.weights = networkWeights; - networkWeights.dispose(); - this.#previousRoundWeights.dispose(); - this.#previousRoundWeights = undefined; + roundValidationDataset, + ), + ); } } @@ -218,36 +205,20 @@ export class Trainer { }, }; + const roundValidationDataset = this.#shouldValidateRound(round) + ? validationDataset + : undefined; + yield this.#runIterationRound( prefixedIterator, this.#roundIterations, - this.#shouldValidateRound(round) ? validationDataset : undefined, + roundValidationDataset, (roundDone) => done = roundDone, - ); - - let roundWeights = this.model.weights; - - if (this.#privacy !== undefined){ - const roundUpdate = roundWeights.sub(this.#previousRoundWeights); - const updateNorm = await Promise.all( - roundUpdate.weights.map(privacy.frobeniusNorm) - ); - this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); - - roundWeights = await applyOptimalPrivacy( - this.#previousRoundWeights, - roundWeights, - this.#privacy, - this.#weightNormHistory, + async () => this.#finishRoundCommunication( Number.MAX_SAFE_INTEGER, - ) - } - - const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); - this.model.weights = networkWeights; - networkWeights.dispose(); - this.#previousRoundWeights.dispose(); - this.#previousRoundWeights = undefined; + roundValidationDataset, + ), + ); round++; if (done) break; @@ -260,6 +231,7 @@ export class Trainer { async *#runRound( dataset: Dataset>, validationDataset?: Dataset>, + onAfterTraining?: () => Promise, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); @@ -277,10 +249,13 @@ export class Trainer { epochsLogs = epochsLogs.push(await epochLogs); } + const postAggregationValidation = await onAfterTraining?.(); + return { epochs: epochsLogs, participants: this.#client.nbOfParticipants, preRoundValidation: validation, + postAggregationValidation, }; } @@ -289,6 +264,7 @@ export class Trainer { maxBatchCount: number, validationDataset?: Dataset>, setDone?: (done: boolean) => void, + onAfterTraining?: () => Promise, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); @@ -313,10 +289,13 @@ export class Trainer { const epochLogs = await result; epochsLogs = epochsLogs.push(epochLogs); + const postAggregationValidation = await onAfterTraining?.(); + return { epochs: epochsLogs, participants: this.#client.nbOfParticipants, preRoundValidation: validation, + postAggregationValidation, }; } @@ -325,6 +304,46 @@ export class Trainer { if (this.#validationFrequency === 0) return false; return round % this.#validationFrequency === 0; } + + async #finishRoundCommunication( + totalRound: number, + validationDataset?: Dataset>, + ): Promise { + let roundWeights = this.model.weights; + + try { + if (this.#privacy !== undefined){ + if (this.#previousRoundWeights === undefined) + throw new Error("previous round weights were not captured"); + + const previousRoundWeights = this.#previousRoundWeights; + const roundUpdate = roundWeights.sub(previousRoundWeights); + const updateNorm = await Promise.all( + roundUpdate.weights.map(privacy.frobeniusNorm) + ); + this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); + + roundWeights = await applyOptimalPrivacy( + previousRoundWeights, + roundWeights, + this.#privacy, + this.#weightNormHistory, + totalRound, + ) + } + + const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); + this.model.weights = networkWeights; + networkWeights.dispose(); + + return this.#shouldValidateAfterAggregation && validationDataset !== undefined + ? await this.model.evaluate(validationDataset) + : undefined; + } finally { + this.#previousRoundWeights?.dispose(); + this.#previousRoundWeights = undefined; + } + } } /** ALDP-FL implementation */ From 3e05567e07c6ad663ade0916d819bc17ee263fcc Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 2 Jun 2026 16:04:23 +0200 Subject: [PATCH 33/89] fix eval script --- cli/src/evaluate_finetuned_gpt2.ts | 170 +++++++++++++++++++++++------ discojs/src/training/disco.ts | 10 +- 2 files changed, 145 insertions(+), 35 deletions(-) diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 84f540b11..490a07499 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -10,33 +10,66 @@ interface Args { testPath: string; maxSamples?: number; savePath?: string; + compareFormats?: boolean; + promptFormat?: PromptFormatName; + contextLength?: number; help?: boolean; } // HOW TO RUN // npm -w cli run eval_finetuned_gpt2 -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/train_no_exp.txt --maxSamples 100 -async function loadDataset(filePath: string, limit = -1): Promise { - const text = await fs.readFile(filePath, "utf-8"); - const lines = text.split("\n"); - - const samples: string[] = []; - let current = ""; - - for (const line of lines) { - const l = line.trim(); - - if (l.includes("<|startoftext|>")) { - current = ""; - } else if (l.includes("<|endoftext|>")) { - samples.push(current.trim()); - if (limit !== -1 && samples.length >= limit) break; - } else { - current += l + "\n"; - } +const PromptFormatNames = [ + "answer-colon-space", + "answer-colon", + "answer-newline", +] as const; + +type PromptFormatName = typeof PromptFormatNames[number]; + +type PromptFormat = { + name: PromptFormatName; + makePrompt: (basePrompt: string) => string; + makeContinuation: (option: string) => string; +}; + +const promptFormats: PromptFormat[] = [ + { + name: "answer-colon-space", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, + makeContinuation: (option) => option, + }, + { + name: "answer-colon", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, + makeContinuation: (option) => ` ${option}`, + }, + { + name: "answer-newline", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, + makeContinuation: (option) => option, + }, +]; + +function castPromptFormatName(raw: string): PromptFormatName { + for (const name of PromptFormatNames) { + if (raw === name) return name; } + throw new Error(`Invalid promptFormat: ${raw}`); +} - return samples; +async function loadDataset(filePath: string, limit = -1): Promise { + const text = await fs.readFile(filePath, "utf-8"); + const samples = text + .split("<|endoftext|>") + .map((sample) => + sample + .replaceAll("<|startoftext|>", "") + .trim(), + ) + .filter((sample) => sample !== ""); + + return limit === -1 ? samples : samples.slice(0, limit); } function parseSample(sample: string) { @@ -46,8 +79,9 @@ function parseSample(sample: string) { const promptLines: string[] = []; for (const line of lines) { - if (line.startsWith("Answer:")) { - answer = line.replace("Answer:", "").trim().charAt(0).toUpperCase(); + const trimmed = line.trim(); + if (trimmed.startsWith("Answer:")) { + answer = trimmed.replace("Answer:", "").trim().charAt(0).toUpperCase(); } else { promptLines.push(line); } @@ -61,13 +95,31 @@ async function scoreContinuation( tfModel: tf.LayersModel, tokenizer: Tokenizer, prompt: string, - continuation: string -): Promise { + continuation: string, + contextLength: number +): Promise<{ score: number; promptTokens: number; continuationTokens: number; usedInputTokens: number }> { const promptTokens = tokenizer.tokenize(prompt).toArray(); const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); + const continuationTokens = fullTokens.length - promptTokens.length; const inputTokens = fullTokens.slice(0, -1); - const inputTensor = tf.tensor2d([inputTokens], [1, inputTokens.length], "int32"); + const offset = Math.max(0, inputTokens.length - contextLength); + const truncatedInputTokens = inputTokens.slice(offset); + + if (truncatedInputTokens.length === 0 || continuationTokens <= 0) { + return { + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens, + usedInputTokens: truncatedInputTokens.length, + }; + } + + const inputTensor = tf.tensor2d( + [truncatedInputTokens], + [1, truncatedInputTokens.length], + "int32", + ); const logits = tfModel.predict(inputTensor) as tf.Tensor; const logProbs = tf.logSoftmax(logits, -1); @@ -78,7 +130,9 @@ async function scoreContinuation( for (let targetPos = promptTokens.length; targetPos < fullTokens.length; targetPos++) { const targetToken = fullTokens[targetPos]; - score += arr[0][targetPos - 1][targetToken]; + const logitPos = targetPos - 1 - offset; + if (logitPos < 0 || logitPos >= arr[0].length) continue; + score += arr[0][logitPos][targetToken]; count++; } @@ -86,16 +140,24 @@ async function scoreContinuation( logits.dispose(); logProbs.dispose(); - return score / count; + return { + score: count === 0 ? Number.NEGATIVE_INFINITY : score / count, + promptTokens: promptTokens.length, + continuationTokens, + usedInputTokens: truncatedInputTokens.length, + }; } async function benchmarkQA( model: models.GPT, tokenizer: Tokenizer, dataset: string[], + format: PromptFormat, + contextLength: number, savePath?: string -) { - console.log("=== QA LOGPROB BENCHMARK ==="); +): Promise { + console.log(`=== QA LOGPROB BENCHMARK (${format.name}) ===`); + console.log(`Context length: ${contextLength}`); const tfModel = model.extract(); @@ -116,6 +178,9 @@ async function benchmarkQA( answer: string; correct: boolean; scores: Record; + promptTokens: number; + continuationTokens: Record; + usedInputTokens: Record; }; const logs: PredictionLog[] = []; @@ -130,11 +195,22 @@ async function benchmarkQA( continue; } - const prompt = `${basePrompt}\nAnswer:`; + const prompt = format.makePrompt(basePrompt); const scores: number[] = []; + const continuationTokens: number[] = []; + const usedInputTokens: number[] = []; for (const opt of options) { - scores.push(await scoreContinuation(tfModel, tokenizer, prompt, ` ${opt}`)); + const result = await scoreContinuation( + tfModel, + tokenizer, + prompt, + format.makeContinuation(opt), + contextLength, + ); + scores.push(result.score); + continuationTokens.push(result.continuationTokens); + usedInputTokens.push(result.usedInputTokens); } let bestIdx = 0; @@ -159,7 +235,14 @@ async function benchmarkQA( predicted, answer, correct: predicted === answer, - scores: scoreMap + scores: scoreMap, + promptTokens: tokenizer.tokenize(prompt).size, + continuationTokens: Object.fromEntries( + options.map((opt, i) => [opt, continuationTokens[i]]) + ), + usedInputTokens: Object.fromEntries( + options.map((opt, i) => [opt, usedInputTokens[i]]) + ), }); if (total % 50 === 0) { @@ -191,6 +274,8 @@ async function benchmarkQA( await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); console.log(`Saved results to ${savePath}`); } + + return accuracy; } async function main() { @@ -199,6 +284,13 @@ async function main() { testPath: { type: String }, maxSamples: { type: Number, optional: true, defaultValue: 100 }, savePath: { type: String, optional: true }, + compareFormats: { type: Boolean, optional: true, defaultValue: false }, + promptFormat: { + type: (raw: string) => castPromptFormatName(raw), + optional: true, + defaultValue: "answer-colon-space", + }, + contextLength: { type: Number, optional: true }, help: { type: Boolean, optional: true } }); @@ -217,9 +309,21 @@ async function main() { console.log(`Loaded ${dataset.length} samples`); - await benchmarkQA(model, tokenizer, dataset, args.savePath); + const contextLength = args.contextLength ?? model.config.contextLength; + const formats = args.compareFormats + ? promptFormats + : promptFormats.filter((format) => format.name === args.promptFormat); + + for (const format of formats) { + const savePath = + args.savePath === undefined || formats.length === 1 + ? args.savePath + : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + + await benchmarkQA(model, tokenizer, dataset, format, contextLength, savePath); + } console.log("Done."); } -main().catch(console.error); \ No newline at end of file +main().catch(console.error); diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 2e2629883..d66782528 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -263,8 +263,6 @@ export class Disco extends EventEmitter<{ `Round: ${roundNum}`, `Initial round loss: ${roundLogs.preRoundValidation?.loss}`, `Initial round accuracy: ${roundLogs.preRoundValidation?.accuracy}`, - `Post-aggregation loss: ${roundLogs.postAggregationValidation?.loss}`, - `Post-aggregation accuracy: ${roundLogs.postAggregationValidation?.accuracy}`, ].join("\n"), ); @@ -286,6 +284,14 @@ export class Disco extends EventEmitter<{ ); } + this.#logger.success( + [ + `Round: ${roundNum}`, + `Post-aggregation loss: ${roundLogs.postAggregationValidation?.loss}`, + `Post-aggregation accuracy: ${roundLogs.postAggregationValidation?.accuracy}`, + ].join("\n"), + ); + return roundLogs; }.bind(this)(); } From 6374b5490b861e383b2c1b7acfd62d90ea24217a Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 2 Jun 2026 19:14:18 +0200 Subject: [PATCH 34/89] add save-checkpoints flag --- cli/src/args.ts | 2 ++ cli/src/cli.ts | 20 ++++++++++++++++++++ cli/src/evaluate_finetuned_gpt2.ts | 15 +++++++++++++-- 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index ec492fbfd..34803996a 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -46,6 +46,7 @@ export interface BenchmarkArguments { saveLogs: boolean saveModel: boolean + saveCheckpoints: boolean host: URL } @@ -78,6 +79,7 @@ const unsafeArgs = parse( goldfishPadTokenId: { type: Number, description: 'Optional padding token id to exclude from Goldfish loss denominator', 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", diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 9cd89a8d6..94f6bf4ad 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -14,6 +14,7 @@ import type { Task, TaskProvider, Network, + Model, } from "@epfml/discojs"; import { Disco, aggregator as aggregators, client as clients } from '@epfml/discojs' @@ -29,6 +30,18 @@ function getOutputDir(): string { return args.outputPath ?? path.join(".", `${args.testID}`); } +async function saveClientModelCheckpoint( + model: Model, + userIndex: number, + round: number, +): Promise { + const checkpointDir = path.join(getOutputDir(), "checkpoints", `round_${round}`); + const checkpointFileName = `client${userIndex}_model.json`; + + await saveModelToDisk(model, checkpointDir, checkpointFileName); + console.log(`Checkpoint saved for client ${userIndex} round ${round} at ${checkpointDir}/${checkpointFileName}`); +} + async function runUser( task: Task, provider: TaskProvider, @@ -76,12 +89,19 @@ async function runUser( try{ debug(`Starting training for client ${userIndex}`); const trainStart = Date.now(); + let lastCheckpointRound: number | undefined = undefined; + for await (const log of disco.trainSummary(data, validationData)){ finalLog.push(log); if (jsonStream){ jsonStream.write(JSON.stringify(log) + "\n"); } + + if (args.saveCheckpoints && lastCheckpointRound !== log.round) { + await saveClientModelCheckpoint(disco.trainer.model, userIndex, log.round); + lastCheckpointRound = log.round; + } } debug(`Training took ${Date.now() - trainStart}ms for client ${userIndex}`); diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 490a07499..59c24de76 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -58,6 +58,16 @@ function castPromptFormatName(raw: string): PromptFormatName { throw new Error(`Invalid promptFormat: ${raw}`); } +function commonPrefixLength(left: number[], right: number[]): number { + const maxLength = Math.min(left.length, right.length); + + for (let i = 0; i < maxLength; i++) { + if (left[i] !== right[i]) return i; + } + + return maxLength; +} + async function loadDataset(filePath: string, limit = -1): Promise { const text = await fs.readFile(filePath, "utf-8"); const samples = text @@ -100,7 +110,8 @@ async function scoreContinuation( ): Promise<{ score: number; promptTokens: number; continuationTokens: number; usedInputTokens: number }> { const promptTokens = tokenizer.tokenize(prompt).toArray(); const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); - const continuationTokens = fullTokens.length - promptTokens.length; + const continuationStart = commonPrefixLength(promptTokens, fullTokens); + const continuationTokens = fullTokens.length - continuationStart; const inputTokens = fullTokens.slice(0, -1); const offset = Math.max(0, inputTokens.length - contextLength); @@ -128,7 +139,7 @@ async function scoreContinuation( let score = 0; let count = 0; - for (let targetPos = promptTokens.length; targetPos < fullTokens.length; targetPos++) { + for (let targetPos = continuationStart; targetPos < fullTokens.length; targetPos++) { const targetToken = fullTokens[targetPos]; const logitPos = targetPos - 1 - offset; if (logitPos < 0 || logitPos >= arr[0].length) continue; From 7a815a922f793f4ff6539802ab7ccbbbb5e3efc5 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 3 Jun 2026 00:45:37 +0200 Subject: [PATCH 35/89] decrease lr --- discojs/src/models/gpt/config.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index c4e68bbe6..114ed8fe0 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -36,7 +36,7 @@ export type GoldfishLossConfig = { } // for a benchmark of performance, see https://github.com/epfml/disco/pull/659 export const DefaultGPTConfig: Required = { - lr: 0.001, + lr: 0.0001, weightDecay: 0, // By default, iterate through the whole dataset and let dataset exhaustion stop the epoch. maxIter: Number.MAX_SAFE_INTEGER, From 718027c743ee421d15886674f83a9b1dd7991f33 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 3 Jun 2026 01:01:57 +0200 Subject: [PATCH 36/89] decrease lr with flag --- cli/src/args.ts | 20 ++++++++++++++++++-- discojs/src/models/gpt/index.ts | 4 ++++ discojs/src/models/gpt/model.ts | 8 ++++++++ 3 files changed, 30 insertions(+), 2 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index 34803996a..567e33b46 100644 --- a/cli/src/args.ts +++ b/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"; @@ -30,6 +30,7 @@ export interface BenchmarkArguments { goldfishK: number goldfishH: number goldfishPadTokenId?: number + learningRate?: number // DP epsilon?: number @@ -77,6 +78,7 @@ const unsafeArgs = parse( 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 }, @@ -235,6 +237,20 @@ 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; + }, }, }; diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 3d480d2c0..808d77d12 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -52,6 +52,10 @@ export class GPT extends Model<"text"> { this.model.setGoldfishLoss(config); } + setLearningRate(lr: number): void { + this.model.setLearningRate(lr); + } + /** * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs) * This allows for getting logs and stopping training without callbacks. diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index cb8571b21..87ea28d1f 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -69,6 +69,14 @@ export class GPTModel extends tf.LayersModel { : tf.train.adam(this.config.lr) } + setLearningRate(lr: number): void { + this.config.lr = lr; + this.optimizer?.dispose(); + this.optimizer = this.config.weightDecay !== 0 + ? getCustomAdam(this, this.config.lr, this.config.weightDecay) + : tf.train.adam(this.config.lr); + } + override async fitDataset(dataset: Dataset, trainingArgs: tf.ModelFitDatasetArgs & { iterationOffset?: number }): Promise { const callbacks = trainingArgs.callbacks as tf.CustomCallbackArgs const evalDataset = trainingArgs.validationData as tf.data.Dataset<{ xs: tf.Tensor2D, ys: tf.Tensor3D }> From f8293470d883baa1a0912bd136cf08d6586a2940 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 4 Jun 2026 13:33:46 +0200 Subject: [PATCH 37/89] fix lr --- cli/src/args.ts | 9 +++++++++ discojs/src/task/training_information.ts | 1 + discojs/src/training/disco.ts | 5 +++++ 3 files changed, 15 insertions(+) diff --git a/cli/src/args.ts b/cli/src/args.ts index 567e33b46..0b52a28de 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -175,6 +175,15 @@ export const args: BenchmarkArguments = { }; } + 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; // For aggregators diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index c9c74dc6d..56c819e36 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -105,6 +105,7 @@ export namespace TrainingInformation { padTokenId: z.number().int().optional(), }) .optional(), + learningRate: z.number().positive().optional(), }), } satisfies Record; diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index d66782528..8df41d062 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -320,9 +320,14 @@ export class Disco extends EventEmitter<{ const configurableModel = model as Model & { setGoldfishLoss?: (config: GoldfishLossConfig | undefined) => void; + setLearningRate?: (learningRate: number) => void; }; configurableModel.setGoldfishLoss?.(this.#task.trainingInformation.goldfishLoss); + if (this.#task.trainingInformation.learningRate !== undefined) { + configurableModel.setLearningRate?.(this.#task.trainingInformation.learningRate); + console.log(`Using GPT learning rate ${this.#task.trainingInformation.learningRate}`); + } } async #preprocessSplitAndBatch( From 7e16a96806d3664d2ba99282e44a23c94944c960 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 4 Jun 2026 19:56:16 +0200 Subject: [PATCH 38/89] set doSample true --- discojs/src/models/gpt/config.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index 114ed8fe0..e164d114b 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -98,7 +98,7 @@ export interface GenerationConfig { export const DefaultGenerationConfig: Required = { temperature: 1.0, - doSample: false, + doSample: true, seed: Math.random(), topk: 50 } From aa5d8be7bc8a2a9f1e91d92b2d96b6e0bb3ad8fe Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 4 Jun 2026 21:47:57 +0200 Subject: [PATCH 39/89] optimize eval script --- cli/src/cli.ts | 4 +- cli/src/evaluate_finetuned_gpt2.ts | 226 +++++++++++++++++++++++------ 2 files changed, 183 insertions(+), 47 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 94f6bf4ad..e96692d1d 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -117,7 +117,9 @@ async function runUser( if (args.saveLogs) { const finalPath = path.join(dir, `client${userIndex}_local_log.json`); - const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, client.ownId, finalLog); + const clientId = + trainingScheme === "local" ? `local-client-${userIndex}` : client.ownId; + const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, clientId, finalLog); await fs.writeFile(finalPath, JSON.stringify(userLog, null, 2)); } diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 59c24de76..00d93c03a 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -68,6 +68,18 @@ function commonPrefixLength(left: number[], right: number[]): number { return maxLength; } +function predictTokenLogits(tfModel: tf.LayersModel, inputTensor: tf.Tensor2D): tf.Tensor3D { + const logits = tfModel.predict(inputTensor); + if (Array.isArray(logits)) { + throw new Error("Expected GPT model to return a single logits tensor"); + } + if (logits.rank !== 3) { + logits.dispose(); + throw new Error(`Expected GPT logits to have rank 3, got rank ${logits.rank}`); + } + return logits as tf.Tensor3D; +} + async function loadDataset(filePath: string, limit = -1): Promise { const text = await fs.readFile(filePath, "utf-8"); const samples = text @@ -101,62 +113,182 @@ function parseSample(sample: string) { return { basePrompt, answer }; } -async function scoreContinuation( +async function scoreContinuations( tfModel: tf.LayersModel, tokenizer: Tokenizer, prompt: string, - continuation: string, + continuations: string[], contextLength: number -): Promise<{ score: number; promptTokens: number; continuationTokens: number; usedInputTokens: number }> { +): Promise<{ score: number; promptTokens: number; continuationTokens: number; usedInputTokens: number }[]> { const promptTokens = tokenizer.tokenize(prompt).toArray(); - const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); - const continuationStart = commonPrefixLength(promptTokens, fullTokens); - const continuationTokens = fullTokens.length - continuationStart; + const scoredInputs = continuations.map((continuation) => { + const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); + const continuationStart = commonPrefixLength(promptTokens, fullTokens); + const continuationTokens = fullTokens.length - continuationStart; + const inputTokens = fullTokens.slice(0, -1); + const offset = Math.max(0, inputTokens.length - contextLength); + const truncatedInputTokens = inputTokens.slice(offset); - const inputTokens = fullTokens.slice(0, -1); - const offset = Math.max(0, inputTokens.length - contextLength); - const truncatedInputTokens = inputTokens.slice(offset); - - if (truncatedInputTokens.length === 0 || continuationTokens <= 0) { return { - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, + fullTokens, + continuationStart, continuationTokens, - usedInputTokens: truncatedInputTokens.length, + offset, + truncatedInputTokens, }; + }); + + const maxInputLength = scoredInputs.reduce( + (maxLength, scoredInput) => Math.max(maxLength, scoredInput.truncatedInputTokens.length), + 0, + ); + + if (maxInputLength === 0) { + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); } + const canScoreFromPromptOnly = scoredInputs.every( + ({ continuationStart, continuationTokens }) => + continuationStart === promptTokens.length && continuationTokens === 1, + ); + + if (canScoreFromPromptOnly) { + const offset = Math.max(0, promptTokens.length - contextLength); + const truncatedPromptTokens = promptTokens.slice(offset); + + if (truncatedPromptTokens.length === 0) { + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: truncatedPromptTokens.length, + })); + } + + const inputTensor = tf.tensor2d( + [truncatedPromptTokens], + [1, truncatedPromptTokens.length], + "int32", + ); + + const optionScores = tf.tidy(() => { + const logits = predictTokenLogits(tfModel, inputTensor); + const lastLogits = logits + .slice([0, truncatedPromptTokens.length - 1, 0], [1, 1, -1]) + .reshape([-1]); + const logProbs = tf.logSoftmax(lastLogits); + const continuationTokenIds = scoredInputs.map( + ({ fullTokens, continuationStart }) => fullTokens[continuationStart], + ); + return tf.gather(logProbs, continuationTokenIds); + }); + + const scores = await optionScores.array() as number[]; + + inputTensor.dispose(); + optionScores.dispose(); + + return scoredInputs.map((scoredInput, index) => ({ + score: scores[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: truncatedPromptTokens.length, + })); + } + + const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ + ...truncatedInputTokens, + ...Array(maxInputLength - truncatedInputTokens.length).fill(0), + ]); + const inputTensor = tf.tensor2d( - [truncatedInputTokens], - [1, truncatedInputTokens.length], + paddedInputs, + [paddedInputs.length, maxInputLength], "int32", ); - const logits = tfModel.predict(inputTensor) as tf.Tensor; - const logProbs = tf.logSoftmax(logits, -1); - const arr = await logProbs.array() as number[][][]; - - let score = 0; - let count = 0; + const targetIndexes: number[][] = []; + const targetTokenIds: number[] = []; + const targetOwners: number[] = []; + + scoredInputs.forEach((scoredInput, batchIdx) => { + const { + fullTokens, + continuationStart, + offset, + truncatedInputTokens, + } = scoredInput; + + // Usually A/B/C/D is one token and the prompt-only fast path above handles it. + // Keep this fallback for prompt/continuation tokenizer merges and multi-token labels. + for (let targetPos = continuationStart; targetPos < fullTokens.length; targetPos++) { + const targetToken = fullTokens[targetPos]; + const logitPos = targetPos - 1 - offset; + if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; + targetIndexes.push([batchIdx, logitPos]); + targetTokenIds.push(targetToken); + targetOwners.push(batchIdx); + } + }); - for (let targetPos = continuationStart; targetPos < fullTokens.length; targetPos++) { - const targetToken = fullTokens[targetPos]; - const logitPos = targetPos - 1 - offset; - if (logitPos < 0 || logitPos >= arr[0].length) continue; - score += arr[0][logitPos][targetToken]; - count++; + if (targetIndexes.length === 0) { + inputTensor.dispose(); + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); } + const logits = predictTokenLogits(tfModel, inputTensor); + const targetLogProbs = tf.tidy(() => { + const targetIndexTensor = tf.tensor2d( + targetIndexes, + [targetIndexes.length, 2], + "int32", + ); + const targetTokenIndexTensor = tf.tensor2d( + targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), + [targetTokenIds.length, 2], + "int32", + ); + const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; + const logProbs = tf.logSoftmax(targetLogits, -1); + return tf.gatherND(logProbs, targetTokenIndexTensor); + }); + + const targetScores = await targetLogProbs.array() as number[]; + const scoreSums = Array(scoredInputs.length).fill(0) as number[]; + const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; + + targetScores.forEach((score, index) => { + const owner = targetOwners[index]; + scoreSums[owner] += score; + scoreCounts[owner]++; + }); + + const results = scoredInputs.map((scoredInput, index) => { + return { + score: scoreCounts[index] === 0 + ? Number.NEGATIVE_INFINITY + : scoreSums[index] / scoreCounts[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + }; + }); + inputTensor.dispose(); logits.dispose(); - logProbs.dispose(); + targetLogProbs.dispose(); - return { - score: count === 0 ? Number.NEGATIVE_INFINITY : score / count, - promptTokens: promptTokens.length, - continuationTokens, - usedInputTokens: truncatedInputTokens.length, - }; + return results; } async function benchmarkQA( @@ -211,14 +343,15 @@ async function benchmarkQA( const scores: number[] = []; const continuationTokens: number[] = []; const usedInputTokens: number[] = []; - for (const opt of options) { - const result = await scoreContinuation( - tfModel, - tokenizer, - prompt, - format.makeContinuation(opt), - contextLength, - ); + const results = await scoreContinuations( + tfModel, + tokenizer, + prompt, + options.map((opt) => format.makeContinuation(opt)), + contextLength, + ); + + for (const result of results) { scores.push(result.score); continuationTokens.push(result.continuationTokens); usedInputTokens.push(result.usedInputTokens); @@ -234,9 +367,10 @@ async function benchmarkQA( if (predicted === answer) correct++; total++; - if (confusion[answer]) { - confusion[answer][predicted]++; + if (confusion[answer]?.[predicted] === undefined) { + throw new Error(`Unexpected confusion matrix key: answer=${answer}, predicted=${predicted}`); } + confusion[answer][predicted]++; const scoreMap = Object.fromEntries( options.map((opt, i) => [opt, scores[i]]) @@ -247,7 +381,7 @@ async function benchmarkQA( answer, correct: predicted === answer, scores: scoreMap, - promptTokens: tokenizer.tokenize(prompt).size, + promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, continuationTokens: Object.fromEntries( options.map((opt, i) => [opt, continuationTokens[i]]) ), From 240c6cdbb6c24441e2161cd7c43600e1c8c131f9 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 5 Jun 2026 01:45:11 +0200 Subject: [PATCH 40/89] add hellaswag-like eval for medMCQ gpt2 --- cli/package.json | 1 + .../evaluate_finetuned_gpt2_full_answer.ts | 500 ++++++++++++++++++ 2 files changed, 501 insertions(+) create mode 100644 cli/src/evaluate_finetuned_gpt2_full_answer.ts diff --git a/cli/package.json b/cli/package.json index 6c4f10bb3..7ea50193f 100644 --- a/cli/package.json +++ b/cli/package.json @@ -10,6 +10,7 @@ "train_gpt": "npm run build && node dist/train_gpt.js", "hellaswag_gpt": "npm run build && node dist/hellaswag_gpt.js", "eval_finetuned_gpt2": "npm run build && node dist/evaluate_finetuned_gpt2.js", + "eval_finetuned_gpt2_full_answer": "npm run build && node dist/evaluate_finetuned_gpt2_full_answer.js", "measure_memorization_gpt2": "npm run build && node dist/measure_memorization_gpt2.js", "finetune_gpt": "npm run build && node dist/finetune_gpt.js", "build": "tsc --build", diff --git a/cli/src/evaluate_finetuned_gpt2_full_answer.ts b/cli/src/evaluate_finetuned_gpt2_full_answer.ts new file mode 100644 index 000000000..661ae9af2 --- /dev/null +++ b/cli/src/evaluate_finetuned_gpt2_full_answer.ts @@ -0,0 +1,500 @@ +import "@tensorflow/tfjs-node"; +import * as tf from "@tensorflow/tfjs"; +import fs from "node:fs/promises"; +import { parse } from "ts-command-line-args"; +import { models, Tokenizer } from "@epfml/discojs"; +import { loadModelFromDisk } from "@epfml/discojs-node"; + +interface Args { + modelPath: string; + testPath: string; + maxSamples?: number; + savePath?: string; + compareFormats?: boolean; + promptFormat?: PromptFormatName; + contextLength?: number; + help?: boolean; +} + +// HOW TO RUN +// npm -w cli run eval_finetuned_gpt2_full_answer -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/test.txt --maxSamples 100 + +const PromptFormatNames = [ + "answer-colon-space", + "answer-colon", + "answer-newline", +] as const; + +type PromptFormatName = typeof PromptFormatNames[number]; + +type PromptFormat = { + name: PromptFormatName; + makePrompt: (basePrompt: string) => string; + makeContinuation: (answer: string) => string; +}; + +type Option = { + label: string; + answer: string; +}; + +type ParsedSample = { + basePrompt: string; + answerLabel: string; + answer: string; + options: Option[]; +}; + +type ScoreResult = { + score: number; + promptTokens: number; + continuationTokens: number; + usedInputTokens: number; +}; + +const promptFormats: PromptFormat[] = [ + { + name: "answer-colon-space", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, + makeContinuation: (answer) => answer, + }, + { + name: "answer-colon", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, + makeContinuation: (answer) => ` ${answer}`, + }, + { + name: "answer-newline", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, + makeContinuation: (answer) => answer, + }, +]; + +function castPromptFormatName(raw: string): PromptFormatName { + for (const name of PromptFormatNames) { + if (raw === name) return name; + } + throw new Error(`Invalid promptFormat: ${raw}`); +} + +function commonPrefixLength(left: number[], right: number[]): number { + const maxLength = Math.min(left.length, right.length); + + for (let i = 0; i < maxLength; i++) { + if (left[i] !== right[i]) return i; + } + + return maxLength; +} + +function predictTokenLogits(tfModel: tf.LayersModel, inputTensor: tf.Tensor2D): tf.Tensor3D { + const logits = tfModel.predict(inputTensor); + if (Array.isArray(logits)) { + throw new Error("Expected GPT model to return a single logits tensor"); + } + if (logits.rank !== 3) { + logits.dispose(); + throw new Error(`Expected GPT logits to have rank 3, got rank ${logits.rank}`); + } + return logits as tf.Tensor3D; +} + +async function loadDataset(filePath: string, limit = -1): Promise { + const text = await fs.readFile(filePath, "utf-8"); + const samples = text + .split("<|endoftext|>") + .map((sample) => + sample + .replaceAll("<|startoftext|>", "") + .trim(), + ) + .filter((sample) => sample !== ""); + + return limit === -1 ? samples : samples.slice(0, limit); +} + +function parseAnswerLine(line: string): { label: string; answer?: string } | undefined { + const match = line.trim().match(/^Answer:\s*([A-D])(?:\.\s*(.*))?$/i); + if (match === null) return undefined; + + const label = match[1].toUpperCase(); + const answerText = match[2]?.trim(); + + return { + label, + answer: answerText === undefined || answerText === "" ? undefined : `${label}. ${answerText}`, + }; +} + +function parseOptionLine(line: string): Option | undefined { + const match = line.trim().match(/^([A-D])\.\s*(.+)$/i); + if (match === null) return undefined; + + const label = match[1].toUpperCase(); + return { + label, + answer: `${label}. ${match[2].trim()}`, + }; +} + +function parseSample(sample: string): ParsedSample { + const lines = sample.split("\n"); + + let answerLabel = ""; + let answerFromLine: string | undefined; + const promptLines: string[] = []; + const options: Option[] = []; + + for (const line of lines) { + const answer = parseAnswerLine(line); + if (answer !== undefined) { + answerLabel = answer.label; + answerFromLine = answer.answer; + continue; + } + + const option = parseOptionLine(line); + if (option !== undefined) { + options.push(option); + } + + promptLines.push(line); + } + + const correctOption = options.find((option) => option.label === answerLabel); + if (correctOption === undefined) { + throw new Error(`Could not match answer label ${JSON.stringify(answerLabel)} to an option`); + } + + if (answerFromLine !== undefined && answerFromLine !== correctOption.answer) { + throw new Error( + `Answer line ${JSON.stringify(answerFromLine)} does not match option ${JSON.stringify(correctOption.answer)}`, + ); + } + + const basePrompt = promptLines.join("\n").trim(); + return { + basePrompt, + answerLabel, + answer: correctOption.answer, + options, + }; +} + +function validateOptions(options: Option[], expectedLabels: string[]): boolean { + if (options.length !== expectedLabels.length) return false; + + const labels = options.map((option) => option.label); + return expectedLabels.every((label) => labels.includes(label)) + && new Set(labels).size === expectedLabels.length; +} + +async function scoreContinuations( + tfModel: tf.LayersModel, + tokenizer: Tokenizer, + prompt: string, + continuations: string[], + contextLength: number, +): Promise { + const promptTokens = tokenizer.tokenize(prompt).toArray(); + const scoredInputs = continuations.map((continuation) => { + const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); + const continuationStart = commonPrefixLength(promptTokens, fullTokens); + const continuationTokens = fullTokens.length - continuationStart; + const inputTokens = fullTokens.slice(0, -1); + const offset = Math.max(0, inputTokens.length - contextLength); + const truncatedInputTokens = inputTokens.slice(offset); + + return { + fullTokens, + continuationStart, + continuationTokens, + offset, + truncatedInputTokens, + }; + }); + + const maxInputLength = scoredInputs.reduce( + (maxLength, scoredInput) => Math.max(maxLength, scoredInput.truncatedInputTokens.length), + 0, + ); + + if (maxInputLength === 0) { + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + } + + const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ + ...truncatedInputTokens, + ...Array(maxInputLength - truncatedInputTokens.length).fill(0), + ]); + + const inputTensor = tf.tensor2d( + paddedInputs, + [paddedInputs.length, maxInputLength], + "int32", + ); + + const targetIndexes: number[][] = []; + const targetTokenIds: number[] = []; + const targetOwners: number[] = []; + + scoredInputs.forEach((scoredInput, batchIdx) => { + const { + fullTokens, + continuationStart, + offset, + truncatedInputTokens, + } = scoredInput; + + // Same ranking as HellaSwag's mean continuation cross-entropy: + // maximize mean log-probability instead of minimizing its negative. + for (let targetPos = continuationStart; targetPos < fullTokens.length; targetPos++) { + const targetToken = fullTokens[targetPos]; + const logitPos = targetPos - 1 - offset; + if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; + targetIndexes.push([batchIdx, logitPos]); + targetTokenIds.push(targetToken); + targetOwners.push(batchIdx); + } + }); + + if (targetIndexes.length === 0) { + inputTensor.dispose(); + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + } + + const logits = predictTokenLogits(tfModel, inputTensor); + const targetLogProbs = tf.tidy(() => { + const targetIndexTensor = tf.tensor2d( + targetIndexes, + [targetIndexes.length, 2], + "int32", + ); + const targetTokenIndexTensor = tf.tensor2d( + targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), + [targetTokenIds.length, 2], + "int32", + ); + const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; + const logProbs = tf.logSoftmax(targetLogits, -1); + return tf.gatherND(logProbs, targetTokenIndexTensor); + }); + + const targetScores = await targetLogProbs.array() as number[]; + const scoreSums = Array(scoredInputs.length).fill(0) as number[]; + const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; + + targetScores.forEach((score, index) => { + const owner = targetOwners[index]; + scoreSums[owner] += score; + scoreCounts[owner]++; + }); + + const results = scoredInputs.map((scoredInput, index) => ({ + score: scoreCounts[index] === 0 + ? Number.NEGATIVE_INFINITY + : scoreSums[index] / scoreCounts[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + + inputTensor.dispose(); + logits.dispose(); + targetLogProbs.dispose(); + + return results; +} + +async function benchmarkFullAnswers( + model: models.GPT, + tokenizer: Tokenizer, + dataset: string[], + format: PromptFormat, + contextLength: number, + savePath?: string, +): Promise { + console.log(`=== FULL ANSWER LOGPROB BENCHMARK (${format.name}) ===`); + console.log(`Context length: ${contextLength}`); + + const tfModel = model.extract(); + + let correct = 0; + let total = 0; + const labels = ["A", "B", "C", "D"]; + const confusion: Record> = Object.fromEntries( + labels.map((label) => [ + label, + Object.fromEntries(labels.map((otherLabel) => [otherLabel, 0])), + ]), + ); + + type PredictionLog = { + predicted: string; + predictedAnswer: string; + answer: string; + answerText: string; + correct: boolean; + scores: Record; + promptTokens: number; + continuationTokens: Record; + usedInputTokens: Record; + }; + + const logs: PredictionLog[] = []; + const start = Date.now(); + + for (const sample of dataset) { + let parsed: ParsedSample; + try { + parsed = parseSample(sample); + } catch (error) { + console.log("Invalid sample:", error instanceof Error ? error.message : error); + continue; + } + + if (!validateOptions(parsed.options, labels)) { + console.log("Invalid options:", parsed.options.map((option) => option.label).join(", ")); + continue; + } + + const prompt = format.makePrompt(parsed.basePrompt); + const results = await scoreContinuations( + tfModel, + tokenizer, + prompt, + parsed.options.map((option) => format.makeContinuation(option.answer)), + contextLength, + ); + + const scores = results.map((result) => result.score); + let bestIdx = 0; + for (let i = 1; i < scores.length; i++) { + if (scores[i] > scores[bestIdx]) bestIdx = i; + } + + const predicted = parsed.options[bestIdx]; + if (predicted.label === parsed.answerLabel) correct++; + total++; + + if (confusion[parsed.answerLabel]?.[predicted.label] === undefined) { + throw new Error( + `Unexpected confusion matrix key: answer=${parsed.answerLabel}, predicted=${predicted.label}`, + ); + } + confusion[parsed.answerLabel][predicted.label]++; + + logs.push({ + predicted: predicted.label, + predictedAnswer: predicted.answer, + answer: parsed.answerLabel, + answerText: parsed.answer, + correct: predicted.label === parsed.answerLabel, + scores: Object.fromEntries( + parsed.options.map((option, i) => [option.label, scores[i]]), + ), + promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, + continuationTokens: Object.fromEntries( + parsed.options.map((option, i) => [option.label, results[i].continuationTokens]), + ), + usedInputTokens: Object.fromEntries( + parsed.options.map((option, i) => [option.label, results[i].usedInputTokens]), + ), + }); + + if (total % 50 === 0) { + console.log(`Processed ${total} samples...`); + } + } + + if (total === 0) { + throw new Error("No valid samples were evaluated"); + } + + const accuracy = correct / total; + const duration = ((Date.now() - start) / 1000).toFixed(2); + + console.log("\n========================="); + console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); + console.log(`Time: ${duration}s`); + console.log("=========================\n"); + + console.log("Confusion Matrix:"); + console.table(confusion); + + console.log("\nPer-class accuracy:"); + for (const cls of labels) { + const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); + const correctCls = confusion[cls][cls]; + const acc = totalCls ? (correctCls / totalCls) * 100 : 0; + + console.log(`${cls}: ${acc.toFixed(2)}%`); + } + + if (savePath) { + await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); + console.log(`Saved results to ${savePath}`); + } + + return accuracy; +} + +async function main() { + const args = parse({ + modelPath: { type: String }, + testPath: { type: String }, + maxSamples: { type: Number, optional: true, defaultValue: 100 }, + savePath: { type: String, optional: true }, + compareFormats: { type: Boolean, optional: true, defaultValue: false }, + promptFormat: { + type: (raw: string) => castPromptFormatName(raw), + optional: true, + defaultValue: "answer-colon-space", + }, + contextLength: { type: Number, optional: true }, + help: { type: Boolean, optional: true }, + }); + + console.log("Loading tokenizer..."); + const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + + console.log("Loading model..."); + const model = await loadModelFromDisk(args.modelPath); + + if (!(model instanceof models.GPT)) { + throw new Error("Model must be GPT"); + } + + console.log("Loading dataset..."); + const dataset = await loadDataset(args.testPath, args.maxSamples); + + console.log(`Loaded ${dataset.length} samples`); + + const contextLength = args.contextLength ?? model.config.contextLength; + const formats = args.compareFormats + ? promptFormats + : promptFormats.filter((format) => format.name === args.promptFormat); + + for (const format of formats) { + const savePath = + args.savePath === undefined || formats.length === 1 + ? args.savePath + : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + + await benchmarkFullAnswers(model, tokenizer, dataset, format, contextLength, savePath); + } + + console.log("Done."); +} + +main().catch(console.error); From d2ae092cefb52fd189d0f0f265aea406ec8ca539 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 12 Jun 2026 18:57:27 +0200 Subject: [PATCH 41/89] add top k in mem script and debug fed client --- cli/src/measure_memorization_gpt2.ts | 63 +++++++++++++++++-- .../src/client/federated/federated_client.ts | 36 +++++++++-- discojs/src/training/disco.ts | 7 +++ 3 files changed, 97 insertions(+), 9 deletions(-) diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts index 6b63aa517..145c3a834 100644 --- a/cli/src/measure_memorization_gpt2.ts +++ b/cli/src/measure_memorization_gpt2.ts @@ -13,12 +13,18 @@ interface Args { promptLengths: string; suffixLength: number; bleuThreshold: number; + decodingStrategy: DecodingStrategy; + temperature: number; + topK: number; seed: number; logEvery: number; savePath?: string; help?: boolean; } +const DecodingStrategies = ["top-k", "greedy"] as const; +type DecodingStrategy = typeof DecodingStrategies[number]; + type PromptResult = { recordIndex: number; recordTokenLength: number; @@ -53,6 +59,14 @@ function parseIntegerList(raw: string): number[] { return values; } +function castDecodingStrategy(raw: string): DecodingStrategy { + for (const strategy of DecodingStrategies) { + if (raw === strategy) return strategy; + } + + throw new Error(`Invalid decodingStrategy: ${raw}`); +} + function seededRandom(seed: number): () => number { let state = seed >>> 0; return () => { @@ -149,11 +163,15 @@ function bleu1to4(reference: number[], candidate: number[]): number { return brevityPenalty * geometricMean; } -async function greedyGenerateGPT2( +async function sampleGenerateGPT2( model: models.GPT, inputIds: number[], maxNewTokens: number, maxContextLength: number, + decodingStrategy: DecodingStrategy, + temperature: number, + topK: number, + seed: number, ): Promise { const generated = [...inputIds]; const tfModel = model.extract(); @@ -170,11 +188,26 @@ async function greedyGenerateGPT2( return output as tf.Tensor; }); - console.log("logits shape:", logits.shape); - const nextTokenTensor = tf.tidy(() => { const last = logits.slice([0, modelInput.length - 1, 0], [1, 1, -1]); - return last.squeeze().argMax(); + const scaled = last.squeeze().div(temperature); + const { values: topKLogits, indices: topKTokens } = tf.topk( + scaled, + topK, + ); + + if (decodingStrategy === "greedy") { + return topKTokens.gather(tf.scalar(0, "int32")).squeeze(); + } + + const sampledIndex = tf.multinomial( + topKLogits.expandDims(0), + 1, + seed + i, + false, + ).squeeze(); + + return topKTokens.gather(sampledIndex).squeeze(); }); const nextTokenData = await nextTokenTensor.data(); @@ -226,6 +259,13 @@ async function main() { promptLengths: { type: String, description: "Comma-separated prompt lengths", defaultValue: "10,50,100,200,500" }, suffixLength: { type: Number, description: "Number of suffix tokens to generate and compare", defaultValue: 50 }, bleuThreshold: { type: Number, description: "BLEU threshold for approximate memorization", defaultValue: 0.75 }, + decodingStrategy: { + type: (raw: string) => castDecodingStrategy(raw), + description: "Generation strategy: top-k or greedy", + defaultValue: "top-k", + }, + temperature: { type: Number, description: "Generation temperature used with top-k sampling", defaultValue: 0.8 }, + topK: { type: Number, description: "Number of most likely tokens considered for top-k sampling", defaultValue: 50 }, seed: { type: Number, description: "Random seed for choosing record split positions", defaultValue: 42 }, logEvery: { type: Number, description: "Print progress every N records; set 0 to disable per-record progress logs", defaultValue: 1 }, savePath: { type: String, description: "Optional JSON output path", optional: true }, @@ -243,6 +283,12 @@ async function main() { ); const promptLengths = parseIntegerList(args.promptLengths); + if (!Number.isFinite(args.temperature) || args.temperature <= 0) { + throw new Error("temperature must be a positive finite number"); + } + if (!Number.isInteger(args.topK) || args.topK < 1) { + throw new Error("topK must be a positive integer"); + } const maxPromptLength = Math.max(...promptLengths); const random = seededRandom(args.seed); @@ -337,11 +383,15 @@ async function main() { } const prompt = ids.slice(splitIndex - promptLength, splitIndex); - const generated = await greedyGenerateGPT2( + const generated = await sampleGenerateGPT2( loadedModel, prompt, args.suffixLength, loadedModel.config.contextLength, + args.decodingStrategy, + args.temperature, + args.topK, + args.seed + recordIndex + promptLength, ); const generatedSuffix = generated.slice(prompt.length, prompt.length + args.suffixLength); @@ -412,6 +462,9 @@ async function main() { promptLengths, suffixLength: args.suffixLength, bleuThreshold: args.bleuThreshold, + decodingStrategy: args.decodingStrategy, + temperature: args.temperature, + topK: args.topK, seed: args.seed, logEvery: args.logEvery, modelContextLength: loadedModel.config.contextLength, diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 9caea57f0..2f3ed70c8 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -12,6 +12,18 @@ import * as messages from "./messages.js"; const debug = createDebug("discojs:client:federated"); +function debugProcessMemory(label: string): void { + if (typeof process === "undefined") return; + + const m = process.memoryUsage(); + debug("%s memory: %O", label, { + rssGB: m.rss / 1024 / 1024 / 1024, + heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, + externalGB: m.external / 1024 / 1024 / 1024, + arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, + }); +} + /** * Arbitrary node id assigned to the federated server which we are communicating with. * Indeed, the server acts as a node within the network. In the federated setting described @@ -146,24 +158,40 @@ export class FederatedClient extends Client<"federated"> { .get(SERVER_NODE_ID); if (payloadToServer === undefined) throw new Error("aggregator didn't make a payload for the server"); + + const round = this.aggregator.round; + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before encode`); + const payload = await serialization.weights.encode(payloadToServer); + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after encode`); + debug( + "[%s] encoded payload for round %d byteLength=%d", + shortenId(this.ownId), + round, + payload.byteLength, + ); + const msg: messages.SendPayload = { type: type.SendPayload, - payload: await serialization.weights.encode(payloadToServer), - round: this.aggregator.round, + payload, + round, }; // Need to await the resulting global model right after sending our local contribution // to make sure we don't miss it - debug(`[${shortenId(this.ownId)}] sent its local update to the server for round ${this.aggregator.round}`); + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before send`); this.server.send(msg); - debug(`[${shortenId(this.ownId)}] is waiting for server update for round ${this.aggregator.round + 1}`); + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after send`); + debug(`[${shortenId(this.ownId)}] sent its local update to the server for round ${round}`); + debug(`[${shortenId(this.ownId)}] is waiting for server update for round ${round + 1}`); const { payload: payloadFromServer, round: serverRound, nbOfParticipants } = await waitMessage( this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update this.nbOfParticipants = nbOfParticipants // Save the current participants + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after receive`); const serverResult = serialization.weights.decode(payloadFromServer); + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after decode server payload`); this.aggregator.setRound(serverRound); return serverResult diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 8df41d062..2d94bac49 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -324,6 +324,13 @@ export class Disco extends EventEmitter<{ }; configurableModel.setGoldfishLoss?.(this.#task.trainingInformation.goldfishLoss); + if (this.#task.trainingInformation.goldfishLoss?.enabled === true) { + const { k, h, padTokenId } = this.#task.trainingInformation.goldfishLoss; + console.log( + `Using Goldfish loss with k=${k}, h=${h}` + + (padTokenId === undefined ? "" : `, padTokenId=${padTokenId}`), + ); + } if (this.#task.trainingInformation.learningRate !== undefined) { configurableModel.setLearningRate?.(this.#task.trainingInformation.learningRate); console.log(`Using GPT learning rate ${this.#task.trainingInformation.learningRate}`); From f37f5de78f2178fbebcee2c4b2928b3508649d92 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Fri, 12 Jun 2026 19:29:18 +0200 Subject: [PATCH 42/89] fix mem leak mean aggr and serialization --- discojs/src/aggregator/mean.ts | 26 ++++++++++++++++--- .../src/client/federated/federated_client.ts | 7 ++++- discojs/src/serialization/model.ts | 6 ++++- .../src/controllers/federated_controller.ts | 12 ++++++--- server/src/routes/training_router.ts | 10 +++++-- 5 files changed, 49 insertions(+), 12 deletions(-) diff --git a/discojs/src/aggregator/mean.ts b/discojs/src/aggregator/mean.ts index eb24d370a..9fa128b3e 100644 --- a/discojs/src/aggregator/mean.ts +++ b/discojs/src/aggregator/mean.ts @@ -2,7 +2,6 @@ import type { Map } from "immutable"; import { AggregationStep } from "./aggregator.js"; import { MultiRoundAggregator, ThresholdType } from "./multiround.js"; import type { WeightsContainer, client } from "../index.js"; -import { aggregation } from "../index.js"; /** * Mean aggregator whose aggregation step consists in computing the mean of the received weights. @@ -19,11 +18,16 @@ export class MeanAggregator extends MultiRoundAggregator { } override _add(nodeId: client.NodeID, contribution: WeightsContainer): void { + const previous = this.contributions.getIn([0, nodeId]) as WeightsContainer | undefined; this.log( this.contributions.hasIn([0, nodeId]) ? AggregationStep.UPDATE : AggregationStep.ADD, nodeId, ); - this.contributions = this.contributions.setIn([0, nodeId], contribution); + if (previous !== undefined) previous.dispose(); + this.contributions = this.contributions.setIn( + [0, nodeId], + contribution.map((weight) => weight.clone()), + ); } override aggregate(): WeightsContainer { @@ -32,8 +36,22 @@ export class MeanAggregator extends MultiRoundAggregator { this.log(AggregationStep.AGGREGATE); - const result = aggregation.avg(currentContributions.values()); - return result; + const contributions = Array.from(currentContributions.values()); + let summed = contributions[0]?.map((weight) => weight.clone()); + if (summed === undefined) throw new Error("aggregating without any contribution"); + + try { + for (const contribution of contributions.slice(1)) { + const next = summed.add(contribution); + summed.dispose(); + summed = next; + } + + return summed.map((weight) => weight.div(contributions.length)); + } finally { + summed.dispose(); + contributions.forEach((contribution) => contribution.dispose()); + } } override makePayloads( diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 2f3ed70c8..db9923220 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -101,7 +101,12 @@ export class FederatedClient extends Client<"federated"> { // which indicates whether there are enough participants or not debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants) if (payload != null) { - model.weights = serialization.weights.decode(payload) + const latestWeights = serialization.weights.decode(payload) + try { + model.weights = latestWeights + } finally { + latestWeights.dispose() + } } return model } diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 4df0fab0c..685229e82 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -103,7 +103,11 @@ export async function decode(encoded: Encoded): Promise> { debug("GPT model weights decoded, deserializing model... CONFIG MIGHT BE WRONG") debug("GPT model config: %O", config || "undefined, using default config") - return models.GPT.deserialize({weights, config}) + try { + return models.GPT.deserialize({weights, config}) + } finally { + weights.dispose() + } } default: throw new Error('invalid encoding, model type unrecognized') diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index bef9a1909..6beba0faf 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -130,10 +130,14 @@ export class FederatedController extends TrainingController< ws.send(msgpack.encode(msg)) debug("Aggregated payload sent to client [%s] for round %o", shortId, this.#aggregator.round) }) - // Add the contribution - debug("Adding contribution from client [%s] to aggregator for round %d", shortId, round) - this.#aggregator.add(clientId, weights, round) - debug(`Successfully added contribution from client [%s] for round ${round}`, shortId) + try { + // Add the contribution + debug("Adding contribution from client [%s] to aggregator for round %d", shortId, round) + this.#aggregator.add(clientId, weights, round) + debug(`Successfully added contribution from client [%s] for round ${round}`, shortId) + } finally { + weights.dispose() + } } else { // If the client sent an invalid or outdated contribution // the server answers with the current round and last global model update diff --git a/server/src/routes/training_router.ts b/server/src/routes/training_router.ts index ef8fd815e..1734ada9d 100644 --- a/server/src/routes/training_router.ts +++ b/server/src/routes/training_router.ts @@ -51,8 +51,14 @@ export class TrainingRouter> { // The federated controller takes the initial model weights at initialization // so that it can send it to new clients - const model = serialization.model.decode(encodedModel) - const encodedWeights = await serialization.weights.encode((await model).weights) + const model = await serialization.model.decode(encodedModel) + const weights = model.weights + let encodedWeights: serialization.Encoded + try { + encodedWeights = await serialization.weights.encode(weights) + } finally { + model[Symbol.dispose]() + } taskController = new FederatedController(t, encodedWeights) } else { const t = task as Task From 153d5aaa4f1acd4ed0eb2eab98770430635f387e Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 13 Jun 2026 02:32:38 +0200 Subject: [PATCH 43/89] add mem fix debug logs --- cli/src/cli.ts | 20 +++++ discojs/src/serialization/model.ts | 10 ++- discojs/src/training/disco.ts | 13 +++ .../src/controllers/federated_controller.ts | 80 ++++++++++++++----- 4 files changed, 100 insertions(+), 23 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index e96692d1d..dbbc8645d 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -30,6 +30,16 @@ function getOutputDir(): string { return args.outputPath ?? path.join(".", `${args.testID}`); } +function debugProcessMemory(label: string): void { + const m = process.memoryUsage(); + debug("%s memory: %O", label, { + rssGB: m.rss / 1024 / 1024 / 1024, + heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, + externalGB: m.external / 1024 / 1024 / 1024, + arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, + }); +} + async function saveClientModelCheckpoint( model: Model, userIndex: number, @@ -88,10 +98,12 @@ async function runUser( try{ debug(`Starting training for client ${userIndex}`); + debugProcessMemory(`client ${userIndex} before training`); const trainStart = Date.now(); let lastCheckpointRound: number | undefined = undefined; for await (const log of disco.trainSummary(data, validationData)){ + debugProcessMemory(`client ${userIndex} round ${log.round} before summary bookkeeping`); finalLog.push(log); if (jsonStream){ @@ -99,22 +111,29 @@ async function runUser( } if (args.saveCheckpoints && lastCheckpointRound !== log.round) { + debugProcessMemory(`client ${userIndex} round ${log.round} before checkpoint`); await saveClientModelCheckpoint(disco.trainer.model, userIndex, log.round); + debugProcessMemory(`client ${userIndex} round ${log.round} after checkpoint`); lastCheckpointRound = log.round; } + debugProcessMemory(`client ${userIndex} round ${log.round} after summary bookkeeping`); } debug(`Training took ${Date.now() - trainStart}ms for client ${userIndex}`); + debugProcessMemory(`client ${userIndex} after training loop`); await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish // Save the trained model if requested if (args.saveModel) { + debugProcessMemory(`client ${userIndex} before final model save`); const modelDir = path.join(getOutputDir(), "models"); const modelFileName = `client${userIndex}_model.json`; await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); + debugProcessMemory(`client ${userIndex} after final model save`); console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); } // saving the entire per-user logs if (args.saveLogs) { + debugProcessMemory(`client ${userIndex} before final log save`); const finalPath = path.join(dir, `client${userIndex}_local_log.json`); const clientId = @@ -122,6 +141,7 @@ async function runUser( const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, clientId, finalLog); await fs.writeFile(finalPath, JSON.stringify(userLog, null, 2)); + debugProcessMemory(`client ${userIndex} after final log save`); } return List(finalLog); diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 685229e82..6f219fb58 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -25,9 +25,13 @@ export async function encode(model: Model): Promise { } case model instanceof models.GPT: { const { weights, config } = model.serialize(); - const serializedWeights = await serialization.weights.encode(weights); - debug("GPT model weights serialized"); - return coder.encode([Type.GPT, serializedWeights, config]); + try { + const serializedWeights = await serialization.weights.encode(weights); + debug("GPT model weights serialized"); + return coder.encode([Type.GPT, serializedWeights, config]); + } finally { + weights.dispose(); + } } default: throw new Error("unknown model type"); diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 2d94bac49..1f920a76b 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -28,6 +28,18 @@ import { RoundLogs, Trainer } from "./trainer.js"; const debug = createDebug("discojs:training:disco"); +function debugProcessMemory(label: string): void { + if (typeof process === "undefined") return; + + const m = process.memoryUsage(); + debug("%s memory: %O", label, { + rssGB: m.rss / 1024 / 1024 / 1024, + heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, + externalGB: m.external / 1024 / 1024 / 1024, + arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, + }); +} + interface DiscoConfig { scheme: N; logger: Logger; @@ -198,6 +210,7 @@ export class Disco extends EventEmitter<{ } const roundLogs = await roundLogsPromise; + debugProcessMemory(`round ${roundNum} after round logs resolved`); for (const {epochNum, epochLogs} of epochResults) { yield buildSummaryLog(roundNum, epochNum, roundLogs, epochLogs); diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 6beba0faf..0d87288d8 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -17,15 +17,26 @@ import FederatedMessages = client.federated.messages const debug = createDebug("server:controllers:federated") +function debugProcessMemory(label: string): void { + const m = process.memoryUsage() + debug("%s memory: %O", label, { + rssGB: m.rss / 1024 / 1024 / 1024, + heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, + externalGB: m.external / 1024 / 1024 / 1024, + arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, + }) +} + export class FederatedController extends TrainingController< D, "federated" > { + #pendingUpdateRecipients = new Map() /** * Aggregators for each hosted task. By default the server waits for 100% of the nodes to send their contributions before aggregating the updates */ - #aggregator = new aggregators.MeanAggregator(undefined, 1, 'relative') + #aggregator = this.#makeAggregator() /** * The most up to date global weights. The model weights are already serialized and * can be sent to participants, before starting training, or when joining mid-training @@ -39,11 +50,49 @@ export class FederatedController extends TrainingController< ) { super(task) this.#latestGlobalWeights = this.initialWeights + } - // Save the latest weight updates to be able to send it to new or outdated clients - this.#aggregator.on('aggregation', async (weightUpdate) => { - this.#latestGlobalWeights = await serialization.weights.encode(weightUpdate) + #makeAggregator(): aggregators.MeanAggregator { + const aggregator = new aggregators.MeanAggregator(undefined, 1, 'relative') + + aggregator.on('aggregation', async (weightUpdate) => { + try { + debugProcessMemory(`round ${aggregator.round} before encoding aggregate`) + const payload = await serialization.weights.encode(weightUpdate) + debugProcessMemory(`round ${aggregator.round} after encoding aggregate`) + debug("round %o aggregate payload byteLength=%d", aggregator.round, payload.byteLength) + this.#latestGlobalWeights = payload + + const recipients = this.#pendingUpdateRecipients + this.#pendingUpdateRecipients = new Map() + + recipients.forEach((recipientWs, recipientId) => { + debug( + "Sending global weights for round %o to client [%s]", + aggregator.round, + recipientId.slice(0, 4), + ) + const msg: FederatedMessages.ReceiveServerPayload = { + type: MessageTypes.ReceiveServerPayload, + round: aggregator.round, + payload, + nbOfParticipants: this.connections.size + } + recipientWs.send(msgpack.encode(msg)) + debug( + "Aggregated payload sent to client [%s] for round %o", + recipientId.slice(0, 4), + aggregator.round, + ) + }) + debugProcessMemory(`round ${aggregator.round} after sending aggregate to ${recipients.size} clients`) + } finally { + weightUpdate.dispose() + debugProcessMemory(`round ${aggregator.round} after disposing aggregate tensor`) + } }) + + return aggregator } /** @@ -115,28 +164,17 @@ export class FederatedController extends TrainingController< this.connections.size, ) const weights = serialization.weights.decode(payload) - - // Create a callback to send the aggregated weight to the client - // when enough contributions are received - this.#aggregator.once('aggregation', async (weightUpdate) => { - debug("Sending global weights for round %o to client [%s]", this.#aggregator.round, shortId) - const msg: FederatedMessages.ReceiveServerPayload = { - type: MessageTypes.ReceiveServerPayload, - round: this.#aggregator.round, // send the current round number after aggregation - payload: await serialization.weights.encode(weightUpdate), - nbOfParticipants: this.connections.size - } - debug("Prepared aggregated payload for client [%s] at round %o", shortId, this.#aggregator.round) - ws.send(msgpack.encode(msg)) - debug("Aggregated payload sent to client [%s] for round %o", shortId, this.#aggregator.round) - }) + let added = false try { // Add the contribution debug("Adding contribution from client [%s] to aggregator for round %d", shortId, round) + this.#pendingUpdateRecipients.set(clientId, ws) this.#aggregator.add(clientId, weights, round) + added = true debug(`Successfully added contribution from client [%s] for round ${round}`, shortId) } finally { weights.dispose() + if (!added) this.#pendingUpdateRecipients.delete(clientId) } } else { // If the client sent an invalid or outdated contribution @@ -163,13 +201,15 @@ export class FederatedController extends TrainingController< ws.on('close', () => { // Remove the participant when the websocket is closed this.connections = this.connections.delete(clientId) + this.#pendingUpdateRecipients.delete(clientId) this.#aggregator.removeNode(clientId) debug("client [%s] left", shortId) // Reset the training session when all participants left if (this.connections.size === 0) { debug("All participants left. Resetting the training session") - this.#aggregator = new aggregators.MeanAggregator(undefined, 1, 'relative') + this.#pendingUpdateRecipients.clear() + this.#aggregator = this.#makeAggregator() this.#latestGlobalWeights = this.initialWeights } From b46fd9b5c6c1d1fc6cd0be9fd6e6030144c3068b Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 13 Jun 2026 02:45:18 +0200 Subject: [PATCH 44/89] fix server error --- discojs/src/serialization/model.ts | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 6f219fb58..685229e82 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -25,13 +25,9 @@ export async function encode(model: Model): Promise { } case model instanceof models.GPT: { const { weights, config } = model.serialize(); - try { - const serializedWeights = await serialization.weights.encode(weights); - debug("GPT model weights serialized"); - return coder.encode([Type.GPT, serializedWeights, config]); - } finally { - weights.dispose(); - } + const serializedWeights = await serialization.weights.encode(weights); + debug("GPT model weights serialized"); + return coder.encode([Type.GPT, serializedWeights, config]); } default: throw new Error("unknown model type"); From d73e5ad78a0d039b05b621ca279806af0d74e9fe Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 13 Jun 2026 17:19:57 +0200 Subject: [PATCH 45/89] fix leak in mem --- cli/src/cli.ts | 28 +++++++++++- .../src/client/federated/federated_client.ts | 44 ++++++++++--------- discojs/src/models/gpt/model.ts | 13 ++++++ .../src/controllers/federated_controller.ts | 16 ++++--- 4 files changed, 72 insertions(+), 29 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index dbbc8645d..35991c535 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -26,6 +26,8 @@ import type { UserLogFile } from "./user_log.js"; const debug = createDebug("cli:main"); +let checkpointQueue = Promise.resolve(); + function getOutputDir(): string { return args.outputPath ?? path.join(".", `${args.testID}`); } @@ -40,6 +42,18 @@ function debugProcessMemory(label: string): void { }); } +function runGarbageCollection(label: string): void { + const gc = (globalThis as typeof globalThis & { gc?: () => void }).gc; + if (gc === undefined) { + debug("%s skipped explicit GC because node was not started with --expose-gc", label); + return; + } + + debugProcessMemory(`${label} before explicit GC`); + gc(); + debugProcessMemory(`${label} after explicit GC`); +} + async function saveClientModelCheckpoint( model: Model, userIndex: number, @@ -52,6 +66,16 @@ async function saveClientModelCheckpoint( console.log(`Checkpoint saved for client ${userIndex} round ${round} at ${checkpointDir}/${checkpointFileName}`); } +async function enqueueClientModelCheckpoint( + model: Model, + userIndex: number, + round: number, +): Promise { + const save = checkpointQueue.then(() => saveClientModelCheckpoint(model, userIndex, round)); + checkpointQueue = save.catch(() => undefined); + await save; +} + async function runUser( task: Task, provider: TaskProvider, @@ -112,8 +136,9 @@ async function runUser( if (args.saveCheckpoints && lastCheckpointRound !== log.round) { debugProcessMemory(`client ${userIndex} round ${log.round} before checkpoint`); - await saveClientModelCheckpoint(disco.trainer.model, userIndex, log.round); + await enqueueClientModelCheckpoint(disco.trainer.model, userIndex, log.round); debugProcessMemory(`client ${userIndex} round ${log.round} after checkpoint`); + runGarbageCollection(`client ${userIndex} round ${log.round} checkpoint`); lastCheckpointRound = log.round; } debugProcessMemory(`client ${userIndex} round ${log.round} after summary bookkeeping`); @@ -129,6 +154,7 @@ async function runUser( const modelFileName = `client${userIndex}_model.json`; await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); debugProcessMemory(`client ${userIndex} after final model save`); + runGarbageCollection(`client ${userIndex} final model save`); console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); } // saving the entire per-user logs diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index db9923220..2a5f26205 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -165,27 +165,29 @@ export class FederatedClient extends Client<"federated"> { throw new Error("aggregator didn't make a payload for the server"); const round = this.aggregator.round; - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before encode`); - const payload = await serialization.weights.encode(payloadToServer); - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after encode`); - debug( - "[%s] encoded payload for round %d byteLength=%d", - shortenId(this.ownId), - round, - payload.byteLength, - ); - - const msg: messages.SendPayload = { - type: type.SendPayload, - payload, - round, - }; - - // Need to await the resulting global model right after sending our local contribution - // to make sure we don't miss it - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before send`); - this.server.send(msg); - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after send`); + { + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before encode`); + const payload = await serialization.weights.encode(payloadToServer); + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after encode`); + debug( + "[%s] encoded payload for round %d byteLength=%d", + shortenId(this.ownId), + round, + payload.byteLength, + ); + + const msg: messages.SendPayload = { + type: type.SendPayload, + payload, + round, + }; + + // Need to await the resulting global model right after sending our local contribution + // to make sure we don't miss it + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before send`); + this.server.send(msg); + debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after send`); + } debug(`[${shortenId(this.ownId)}] sent its local update to the server for round ${round}`); debug(`[${shortenId(this.ownId)}] is waiting for server update for round ${round + 1}`); const { diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 87ea28d1f..4a3f6aa70 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -9,6 +9,18 @@ import { GPTArchitecture } from './layers.js' const debug = createDebug("discojs:models:gpt:model"); +function processMemory(): Record | undefined { + if (typeof process === "undefined") return undefined; + + const m = process.memoryUsage(); + return { + rssGB: m.rss / 1024 / 1024 / 1024, + heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, + externalGB: m.external / 1024 / 1024 / 1024, + arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, + }; +} + /** * tfjs does not export LazyIterator and Dataset... */ @@ -170,6 +182,7 @@ export class GPTModel extends tf.LayersModel { loss, memory, allocated: tf.memory().numTensors, + processMemory: processMemory(), preprocessingTime, weightUpdateTime, }); diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 0d87288d8..6dea296e8 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -65,6 +65,14 @@ export class FederatedController extends TrainingController< const recipients = this.#pendingUpdateRecipients this.#pendingUpdateRecipients = new Map() + const msg: FederatedMessages.ReceiveServerPayload = { + type: MessageTypes.ReceiveServerPayload, + round: aggregator.round, + payload, + nbOfParticipants: this.connections.size + } + const encodedMsg = msgpack.encode(msg) + debugProcessMemory(`round ${aggregator.round} after encoding websocket message`) recipients.forEach((recipientWs, recipientId) => { debug( @@ -72,13 +80,7 @@ export class FederatedController extends TrainingController< aggregator.round, recipientId.slice(0, 4), ) - const msg: FederatedMessages.ReceiveServerPayload = { - type: MessageTypes.ReceiveServerPayload, - round: aggregator.round, - payload, - nbOfParticipants: this.connections.size - } - recipientWs.send(msgpack.encode(msg)) + recipientWs.send(encodedMsg) debug( "Aggregated payload sent to client [%s] for round %o", recipientId.slice(0, 4), From 95952cbf877a82a1fbd4e7687b4384fe31bbbf3c Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 13 Jun 2026 17:27:38 +0200 Subject: [PATCH 46/89] patch DP --- discojs/src/training/trainer.ts | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index ed0a72c7e..0d1af4079 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -178,6 +178,15 @@ export class Trainer { if (this.#roundIterations === undefined) throw new Error("roundIterations was not set"); + const totalRound = + this.#privacy?.differentialPrivacy === undefined + ? Number.MAX_SAFE_INTEGER + : Math.max( + 1, + Math.ceil((await dataset.size()) / this.#roundIterations) * + this.#epochs, + ); + let round = 0; for (let epoch = 0; epoch < this.#epochs; epoch++) { const trainingIterator = dataset[Symbol.asyncIterator](); @@ -215,7 +224,7 @@ export class Trainer { roundValidationDataset, (roundDone) => done = roundDone, async () => this.#finishRoundCommunication( - Number.MAX_SAFE_INTEGER, + totalRound, roundValidationDataset, ), ); @@ -404,6 +413,19 @@ async function applyOptimalPrivacy( ) : dpClippingRadius; + const sigmas = effectiveRadius.map((r) => + (2 * r * Math.sqrt(2 * Math.log(1.25 / delta))) / epsilon, + ); + debug("DP applied: %O", { + totalRound, + epsilon, + delta, + radiusMin: Math.min(...effectiveRadius), + radiusMax: Math.max(...effectiveRadius), + sigmaMin: Math.min(...sigmas), + sigmaMax: Math.max(...sigmas), + }); + ret = previousEpochWeights.add( await privacy.addOptimalNoise( weightsProgress, From d46614a80045b1f9e278b7f1620b63f13f78c64e Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 13 Jun 2026 19:12:31 +0200 Subject: [PATCH 47/89] patch DP mem leak --- discojs/src/privacy.ts | 14 +++++-- discojs/src/training/trainer.ts | 69 ++++++++++++++++++++------------- 2 files changed, 53 insertions(+), 30 deletions(-) diff --git a/discojs/src/privacy.ts b/discojs/src/privacy.ts index 8ed51ed1c..70dc6dde4 100644 --- a/discojs/src/privacy.ts +++ b/discojs/src/privacy.ts @@ -6,7 +6,9 @@ import type { WeightNormHistory } from "./training/trainer.js"; /** Computes the Frobenius norm of the given weights. */ export async function frobeniusNorm(weights: tf.Tensor): Promise { - const squared = await weights.square().sum().data(); + const squaredTensor = tf.tidy(() => weights.square().sum()); + const squared = await squaredTensor.data(); + squaredTensor.dispose(); if (squared.length !== 1) throw new Error("unexpected weights shape"); return Math.sqrt(squared[0]); } @@ -47,9 +49,13 @@ export async function addOptimalNoise( const sigmas = sens.map((s)=>(s * Math.sqrt(2*Math.log(1.25/delta))/epsilon)); const clippedWeights = await clipNorm(weightUpdates, clippingRadius); - return clippedWeights.map((w, i) => - w.add(tf.randomNormal(w.shape, 0, sigmas[i])) - ) + try { + return clippedWeights.map((w, i) => + tf.tidy(() => w.add(tf.randomNormal(w.shape, 0, sigmas[i]))) + ) + } finally { + clippedWeights.dispose(); + } } /** diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 0d1af4079..b2a75080f 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -327,18 +327,24 @@ export class Trainer { const previousRoundWeights = this.#previousRoundWeights; const roundUpdate = roundWeights.sub(previousRoundWeights); - const updateNorm = await Promise.all( - roundUpdate.weights.map(privacy.frobeniusNorm) - ); - this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); - - roundWeights = await applyOptimalPrivacy( - previousRoundWeights, - roundWeights, - this.#privacy, - this.#weightNormHistory, - totalRound, - ) + try { + const updateNorm = await Promise.all( + roundUpdate.weights.map(privacy.frobeniusNorm) + ); + this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); + } finally { + roundUpdate.dispose(); + } + + const privateRoundWeights = await applyOptimalPrivacy( + previousRoundWeights, + roundWeights, + this.#privacy, + this.#weightNormHistory, + totalRound, + ); + roundWeights.dispose(); + roundWeights = privateRoundWeights; } const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); @@ -349,6 +355,7 @@ export class Trainer { ? await this.model.evaluate(validationDataset) : undefined; } finally { + roundWeights.dispose(); this.#previousRoundWeights?.dispose(); this.#previousRoundWeights = undefined; } @@ -377,14 +384,19 @@ async function applyOptimalPrivacy( const previousRoundWeights = previous ?? current.map((w) => tf.zerosLike(w)); const weightsProgress = current.sub(previousRoundWeights); - ret = previousRoundWeights.add( - await privacy.clipNorm( - weightsProgress, - Repeat(options.byzantineFaultTolerance.clippingRadius) - .take(weightsProgress.weights.length) - .toArray(), - ), + const clippedProgress = await privacy.clipNorm( + weightsProgress, + Repeat(options.byzantineFaultTolerance.clippingRadius) + .take(weightsProgress.weights.length) + .toArray(), ); + try { + ret = previousRoundWeights.add(clippedProgress); + } finally { + weightsProgress.dispose(); + clippedProgress.dispose(); + if (previous === undefined) previousRoundWeights.dispose(); + } } // Adding Gaussian noise for DP @@ -426,14 +438,19 @@ async function applyOptimalPrivacy( sigmaMax: Math.max(...sigmas), }); - ret = previousEpochWeights.add( - await privacy.addOptimalNoise( - weightsProgress, - epsilon, - delta, - effectiveRadius, - ), + const noisyProgress = await privacy.addOptimalNoise( + weightsProgress, + epsilon, + delta, + effectiveRadius, ); + try { + ret = previousEpochWeights.add(noisyProgress); + } finally { + weightsProgress.dispose(); + noisyProgress.dispose(); + if (previous === undefined) previousEpochWeights.dispose(); + } } return ret; } From b2fbd1e609368a1d552a42535594b429e3c03727 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sat, 13 Jun 2026 19:36:16 +0200 Subject: [PATCH 48/89] fix unwanted dispose --- discojs/src/training/trainer.ts | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index b2a75080f..cf4969194 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -319,6 +319,7 @@ export class Trainer { validationDataset?: Dataset>, ): Promise { let roundWeights = this.model.weights; + let disposeRoundWeightsAfterSend = false; try { if (this.#privacy !== undefined){ @@ -343,8 +344,8 @@ export class Trainer { this.#weightNormHistory, totalRound, ); - roundWeights.dispose(); roundWeights = privateRoundWeights; + disposeRoundWeightsAfterSend = true; } const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); @@ -355,7 +356,7 @@ export class Trainer { ? await this.model.evaluate(validationDataset) : undefined; } finally { - roundWeights.dispose(); + if (disposeRoundWeightsAfterSend) roundWeights.dispose(); this.#previousRoundWeights?.dispose(); this.#previousRoundWeights = undefined; } From 1ed4bee54aa82ad1f7dcd706f7c6c22069d90313 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 6 Jul 2026 18:09:12 +0200 Subject: [PATCH 49/89] clean lint --- cli/src/cli.ts | 39 +------------------------ cli/src/evaluate_finetuned_gpt2.ts | 2 +- cli/src/measure_memorization_gpt2.ts | 5 ++-- discojs/src/default_tasks/privacyrun.ts | 2 -- discojs/src/models/gpt/config.ts | 4 +-- discojs/src/models/gpt/index.ts | 8 +---- discojs/src/models/gpt/model.ts | 1 - 7 files changed, 7 insertions(+), 54 deletions(-) diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 35991c535..d52c36abe 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -32,26 +32,13 @@ function getOutputDir(): string { return args.outputPath ?? path.join(".", `${args.testID}`); } -function debugProcessMemory(label: string): void { - const m = process.memoryUsage(); - debug("%s memory: %O", label, { - rssGB: m.rss / 1024 / 1024 / 1024, - heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, - externalGB: m.external / 1024 / 1024 / 1024, - arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, - }); -} - function runGarbageCollection(label: string): void { const gc = (globalThis as typeof globalThis & { gc?: () => void }).gc; if (gc === undefined) { debug("%s skipped explicit GC because node was not started with --expose-gc", label); return; } - - debugProcessMemory(`${label} before explicit GC`); gc(); - debugProcessMemory(`${label} after explicit GC`); } async function saveClientModelCheckpoint( @@ -78,7 +65,6 @@ async function enqueueClientModelCheckpoint( async function runUser( task: Task, - provider: TaskProvider, url: URL, data: Dataset, validationData: Dataset | undefined, @@ -86,7 +72,6 @@ async function runUser( numberOfUsers: number, ): Promise> { debug(`Starting runUser for client ${userIndex}`); - const userStart = Date.now(); const trainingScheme = task.trainingInformation.scheme as N const aggregator = aggregators.getAggregator(task) const client = clients.getClient(trainingScheme, url, task, aggregator) @@ -95,18 +80,6 @@ async function runUser( preprocessOnce: false, debugLabel: `client${userIndex}`, }); - - // For local training, load model from provider before training starts - // if (trainingScheme === "local") { - // debug(`Loading model for training client ${userIndex}...`); - // const modelStart = Date.now(); - // console.log("Loading model for local training..."); - // disco.trainer.model = await provider.getModel(); - // console.log("Model loaded successfully"); - // debug(`Model loading took ${Date.now() - modelStart}ms for client ${userIndex}`); - // } - - const dir = getOutputDir(); await fs.mkdir(dir, { recursive: true }); @@ -122,12 +95,10 @@ async function runUser( try{ debug(`Starting training for client ${userIndex}`); - debugProcessMemory(`client ${userIndex} before training`); const trainStart = Date.now(); let lastCheckpointRound: number | undefined = undefined; for await (const log of disco.trainSummary(data, validationData)){ - debugProcessMemory(`client ${userIndex} round ${log.round} before summary bookkeeping`); finalLog.push(log); if (jsonStream){ @@ -135,31 +106,24 @@ async function runUser( } if (args.saveCheckpoints && lastCheckpointRound !== log.round) { - debugProcessMemory(`client ${userIndex} round ${log.round} before checkpoint`); await enqueueClientModelCheckpoint(disco.trainer.model, userIndex, log.round); - debugProcessMemory(`client ${userIndex} round ${log.round} after checkpoint`); runGarbageCollection(`client ${userIndex} round ${log.round} checkpoint`); lastCheckpointRound = log.round; } - debugProcessMemory(`client ${userIndex} round ${log.round} after summary bookkeeping`); } debug(`Training took ${Date.now() - trainStart}ms for client ${userIndex}`); - debugProcessMemory(`client ${userIndex} after training loop`); await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish // Save the trained model if requested if (args.saveModel) { - debugProcessMemory(`client ${userIndex} before final model save`); const modelDir = path.join(getOutputDir(), "models"); const modelFileName = `client${userIndex}_model.json`; await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); - debugProcessMemory(`client ${userIndex} after final model save`); runGarbageCollection(`client ${userIndex} final model save`); console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); } // saving the entire per-user logs if (args.saveLogs) { - debugProcessMemory(`client ${userIndex} before final log save`); const finalPath = path.join(dir, `client${userIndex}_local_log.json`); const clientId = @@ -167,7 +131,6 @@ async function runUser( const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, clientId, finalLog); await fs.writeFile(finalPath, JSON.stringify(userLog, null, 2)); - debugProcessMemory(`client ${userIndex} after final log save`); } return List(finalLog); @@ -217,7 +180,7 @@ async function main( } const logs = await Promise.all( - dataSplits.map((data, i) => runUser(task, provider, args.host, data as Dataset, validationData, i, numberOfUsers)) + dataSplits.map((data, i) => runUser(task, args.host, data as Dataset, validationData, i, numberOfUsers)) ) if (args.saveLogs) { diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 00d93c03a..f6f01874c 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -188,7 +188,7 @@ async function scoreContinuations( return tf.gather(logProbs, continuationTokenIds); }); - const scores = await optionScores.array() as number[]; + const scores = await optionScores.array(); inputTensor.dispose(); optionScores.dispose(); diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts index 145c3a834..fbbdb7b38 100644 --- a/cli/src/measure_memorization_gpt2.ts +++ b/cli/src/measure_memorization_gpt2.ts @@ -183,9 +183,9 @@ async function sampleGenerateGPT2( const logits = tf.tidy(() => { const output = tfModel.predict(input); if (Array.isArray(output)) { - return output[0] as tf.Tensor; + return output[0]; } - return output as tf.Tensor; + return output; }); const nextTokenTensor = tf.tidy(() => { @@ -289,7 +289,6 @@ async function main() { if (!Number.isInteger(args.topK) || args.topK < 1) { throw new Error("topK must be a positive integer"); } - const maxPromptLength = Math.max(...promptLengths); const random = seededRandom(args.seed); console.log("Loading tokenizer..."); diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index 04b031564..d1bb23863 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -44,8 +44,6 @@ export const privacyrun: TaskProvider<"text", "federated"> = { // The model should be in DiscoJS serialization format (created by onnx-converter) // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model.json"; - // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; - const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; try { diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index e164d114b..c4e68bbe6 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -36,7 +36,7 @@ export type GoldfishLossConfig = { } // for a benchmark of performance, see https://github.com/epfml/disco/pull/659 export const DefaultGPTConfig: Required = { - lr: 0.0001, + lr: 0.001, weightDecay: 0, // By default, iterate through the whole dataset and let dataset exhaustion stop the epoch. maxIter: Number.MAX_SAFE_INTEGER, @@ -98,7 +98,7 @@ export interface GenerationConfig { export const DefaultGenerationConfig: Required = { temperature: 1.0, - doSample: true, + doSample: false, seed: Math.random(), topk: 50 } diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 808d77d12..e13591ba7 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -277,13 +277,7 @@ export class GPT extends Model<"text"> { debug("GPT model deserialization started") - const config = - data.config === undefined - ? undefined - : { - ...data.config, - maxIter: DefaultGPTConfig.maxIter, - }; + const config = data.config; const model = new GPT(config); diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 4a3f6aa70..45f590c54 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -114,7 +114,6 @@ export class GPTModel extends tf.LayersModel { let preprocessingTime = performance.now() await Promise.all([xs.data(), ys.data()]) - // await Promise.resolve() preprocessingTime = performance.now() - preprocessingTime // TODO include as a tensor inside the model From 2a4aa769e3123452a7aab8e5b92e8f6fc7f5579a Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 6 Jul 2026 20:00:11 +0200 Subject: [PATCH 50/89] failing test fix --- discojs/src/client/federated/federated_client.ts | 9 +-------- server/tests/e2e/federated.spec.ts | 8 ++++++-- 2 files changed, 7 insertions(+), 10 deletions(-) diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 2a5f26205..d3765eb5a 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -122,16 +122,9 @@ export class FederatedClient extends Client<"federated"> { this.aggregator.setNodes(this.aggregator.nodes.delete(SERVER_NODE_ID)); } - // override onRoundBeginCommunication(): Promise { - // // Prepare the result promise for the incoming round - // this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve)) - // this.saveAndEmit("local training") - // return Promise.resolve(); - // } - override async onRoundBeginCommunication(): Promise { + override onRoundBeginCommunication(): Promise { // Prepare the result promise for the incoming round this.aggregationResult = new Promise((resolve) => this.aggregator.once('aggregation', resolve)) - await this.waitForParticipantsIfNeeded() // In case we are waiting for more participants, we wait before starting the local training this.saveAndEmit("local training") return Promise.resolve(); } diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts index 20ba132e2..fcf248227 100644 --- a/server/tests/e2e/federated.spec.ts +++ b/server/tests/e2e/federated.spec.ts @@ -9,7 +9,7 @@ import type { TaskProvider, WeightsContainer, } from "@epfml/discojs"; -import { Disco, defaultTasks } from "@epfml/discojs"; +import { Disco, defaultTasks, models } from "@epfml/discojs"; import { List } from "immutable"; import { assert, afterEach, describe, expect, it } from "vitest"; import { Server } from "../../src/index.js"; @@ -161,6 +161,11 @@ describe("end-to-end federated", () => { }; const url = await startServer({ ...defaultTasks.wikitext, + getModel: () => + Promise.resolve(new models.GPT({ + contextLength: task.trainingInformation.contextLength, + maxIter: 10, + })), getTask: () => Promise.resolve(task), }); const dataset = datasets.loadWikitext(); @@ -353,4 +358,3 @@ describe("end-to-end federated", () => { assert.isTrue(m1.equals(m2) && m2.equals(m3)); }) }); - From cddb271e03c7da747a101f0105abf6bb8be26643 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 6 Jul 2026 20:14:48 +0200 Subject: [PATCH 51/89] trainer fix deadlock --- discojs/src/training/trainer.ts | 25 +++++-------------------- 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index cf4969194..4e75233a0 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -157,14 +157,9 @@ export class Trainer { ? validationDataset : undefined; - yield this.#runRound( - dataset, - roundValidationDataset, - async () => this.#finishRoundCommunication( - totalRound, - roundValidationDataset, - ), - ); + yield this.#runRound(dataset, roundValidationDataset); + + await this.#finishRoundCommunication(totalRound); } } @@ -223,12 +218,10 @@ export class Trainer { this.#roundIterations, roundValidationDataset, (roundDone) => done = roundDone, - async () => this.#finishRoundCommunication( - totalRound, - roundValidationDataset, - ), ); + await this.#finishRoundCommunication(totalRound); + round++; if (done) break; next = await trainingIterator.next(); @@ -240,7 +233,6 @@ export class Trainer { async *#runRound( dataset: Dataset>, validationDataset?: Dataset>, - onAfterTraining?: () => Promise, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); @@ -258,13 +250,10 @@ export class Trainer { epochsLogs = epochsLogs.push(await epochLogs); } - const postAggregationValidation = await onAfterTraining?.(); - return { epochs: epochsLogs, participants: this.#client.nbOfParticipants, preRoundValidation: validation, - postAggregationValidation, }; } @@ -273,7 +262,6 @@ export class Trainer { maxBatchCount: number, validationDataset?: Dataset>, setDone?: (done: boolean) => void, - onAfterTraining?: () => Promise, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); @@ -298,13 +286,10 @@ export class Trainer { const epochLogs = await result; epochsLogs = epochsLogs.push(epochLogs); - const postAggregationValidation = await onAfterTraining?.(); - return { epochs: epochsLogs, participants: this.#client.nbOfParticipants, preRoundValidation: validation, - postAggregationValidation, }; } From 2bec4093f7c61a17bb0d8f36a7305fe49730193e Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Mon, 6 Jul 2026 21:39:13 +0200 Subject: [PATCH 52/89] fix dispose --- discojs/src/client/client.ts | 1 + discojs/src/client/federated/federated_client.ts | 6 +----- discojs/src/serialization/model.ts | 6 +----- discojs/src/training/trainer.ts | 1 - server/src/controllers/federated_controller.ts | 5 +++-- 5 files changed, 6 insertions(+), 13 deletions(-) diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 5c8aa1304..6f2d122ac 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -159,6 +159,7 @@ export abstract class Client extends EventEmitter<{ // Emit the last status emitted before waiting if defined if (this.#previousStatus !== undefined) this.emit("status", this.#previousStatus) this.nbOfParticipants = event.nbOfParticipants + this.promiseForMoreParticipants = undefined resolve() }) }) diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index d3765eb5a..299316b46 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -102,11 +102,7 @@ export class FederatedClient extends Client<"federated"> { debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants) if (payload != null) { const latestWeights = serialization.weights.decode(payload) - try { - model.weights = latestWeights - } finally { - latestWeights.dispose() - } + model.weights = latestWeights } return model } diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 685229e82..4df0fab0c 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -103,11 +103,7 @@ export async function decode(encoded: Encoded): Promise> { debug("GPT model weights decoded, deserializing model... CONFIG MIGHT BE WRONG") debug("GPT model config: %O", config || "undefined, using default config") - try { - return models.GPT.deserialize({weights, config}) - } finally { - weights.dispose() - } + return models.GPT.deserialize({weights, config}) } default: throw new Error('invalid encoding, model type unrecognized') diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 4e75233a0..d4240be4b 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -335,7 +335,6 @@ export class Trainer { const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); this.model.weights = networkWeights; - networkWeights.dispose(); return this.#shouldValidateAfterAggregation && validationDataset !== undefined ? await this.model.evaluate(validationDataset) diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 6dea296e8..de84b84eb 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -57,14 +57,15 @@ export class FederatedController extends TrainingController< aggregator.on('aggregation', async (weightUpdate) => { try { + const recipients = this.#pendingUpdateRecipients + this.#pendingUpdateRecipients = new Map() + debugProcessMemory(`round ${aggregator.round} before encoding aggregate`) const payload = await serialization.weights.encode(weightUpdate) debugProcessMemory(`round ${aggregator.round} after encoding aggregate`) debug("round %o aggregate payload byteLength=%d", aggregator.round, payload.byteLength) this.#latestGlobalWeights = payload - const recipients = this.#pendingUpdateRecipients - this.#pendingUpdateRecipients = new Map() const msg: FederatedMessages.ReceiveServerPayload = { type: MessageTypes.ReceiveServerPayload, round: aggregator.round, From 3061d940c859d57091cd98509ba5aed90aa09ec4 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Tue, 7 Jul 2026 00:55:21 +0200 Subject: [PATCH 53/89] fix debug line --- discojs/src/serialization/model.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 4df0fab0c..94f08baf9 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -98,10 +98,10 @@ export async function decode(encoded: Encoded): Promise> { "invalid encoding, gpt-tfjs model weights should be an encoding of its weights", ); - debug("GPT model weights decoding...") + debug("GPT model weights decoding") const weights = serialization.weights.decode(rawModel) - debug("GPT model weights decoded, deserializing model... CONFIG MIGHT BE WRONG") + debug("GPT model weights decoded") debug("GPT model config: %O", config || "undefined, using default config") return models.GPT.deserialize({weights, config}) } From 4ad761b81778bb5e4a971cedb1e010221250a938 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 8 Jul 2026 15:51:18 +0200 Subject: [PATCH 54/89] continue merge --- cli/package.json | 35 +------ cli/src/args.ts | 76 -------------- cli/src/cli.ts | 64 ------------ cli/src/data.ts | 17 ---- discojs/src/client/client.ts | 9 -- discojs/src/client/event_connection.ts | 16 --- .../src/client/federated/federated_client.ts | 38 ------- discojs/src/client/federated/messages.ts | 9 -- discojs/src/client/local_client.ts | 8 -- discojs/src/default_tasks/index.ts | 10 -- discojs/src/models/gpt/config.ts | 22 ----- discojs/src/models/gpt/index.ts | 5 - discojs/src/models/gpt/model.ts | 99 ------------------- discojs/src/models/index.ts | 6 -- discojs/src/privacy.ts | 12 --- 15 files changed, 4 insertions(+), 422 deletions(-) diff --git a/cli/package.json b/cli/package.json index 635042db1..3c8f9fd54 100644 --- a/cli/package.json +++ b/cli/package.json @@ -1,34 +1,4 @@ { -<<<<<<< HEAD - "name": "cli", - "private": true, - "type": "module", - "main": "dist/cli.js", - "scripts": { - "watch": "nodemon --ext ts --ignore dist --watch ../discojs-node/dist --watch ../server/dist --watch . --exec npm run", - "start": "npm run build && node dist/cli.js", - "benchmark_gpt": "npm run build && node dist/benchmark_gpt.js", - "train_gpt": "npm run build && node dist/train_gpt.js", - "hellaswag_gpt": "npm run build && node dist/hellaswag_gpt.js", - "eval_finetuned_gpt2": "npm run build && node dist/evaluate_finetuned_gpt2.js", - "eval_finetuned_gpt2_full_answer": "npm run build && node dist/evaluate_finetuned_gpt2_full_answer.js", - "measure_memorization_gpt2": "npm run build && node dist/measure_memorization_gpt2.js", - "finetune_gpt": "npm run build && node dist/finetune_gpt.js", - "build": "tsc --build", - "test": ": nothing" - }, - "author": "", - "license": "ISC", - "dependencies": { - "@epfml/discojs-node": "*", - "server": "*", - "tslib": "2" - }, - "devDependencies": { - "nodemon": "3", - "ts-command-line-args": "2" - } -======= "name": "cli", "private": true, "type": "module", @@ -39,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" }, @@ -55,5 +29,4 @@ "nodemon": "3.1.14", "ts-command-line-args": "2.5.1" } ->>>>>>> develop } diff --git a/cli/src/args.ts b/cli/src/args.ts index 6a14b3b12..11aa3982f 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -2,11 +2,7 @@ import { parse } from "ts-command-line-args"; import { Map, Set } from "immutable"; import type { DataType, Network, TaskProvider } from "@epfml/discojs"; -<<<<<<< HEAD import { defaultTasks, models } from '@epfml/discojs' -======= -import { defaultTasks } from "@epfml/discojs"; ->>>>>>> develop type AggregationStrategy = "mean" | "byzantine" | "secure"; @@ -16,7 +12,6 @@ function parseAggregator(raw: string): AggregationStrategy { } export interface BenchmarkArguments { -<<<<<<< HEAD provider: TaskProvider; testID: string numberOfUsers: number @@ -34,15 +29,6 @@ export interface BenchmarkArguments { goldfishH: number goldfishPadTokenId?: number learningRate?: number -======= - provider: TaskProvider; - testID: string; - numberOfUsers: number; - epochs: number; - roundDuration: number; - batchSize: number; - validationSplit: number; ->>>>>>> develop // DP epsilon?: number; @@ -57,7 +43,6 @@ export interface BenchmarkArguments { // Secure aggregator maxShareValue?: number; -<<<<<<< HEAD saveLogs: boolean saveModel: boolean saveCheckpoints: boolean @@ -70,22 +55,11 @@ type BenchmarkUnsafeArguments = Omit & { validationDatasetPath?: string help?: boolean } -======= - save: boolean; - host: URL; -} - -type BenchmarkUnsafeArguments = Omit & { - task: string; - help?: boolean; -}; ->>>>>>> develop const argExample = "e.g. pnpm start -u 2 -e 3 # runs 2 users for 3 epochs"; const unsafeArgs = parse( { -<<<<<<< HEAD testID: { type: String, alias: 'i', description: 'ID of the testcase' }, task: { type: String, alias: 't', description: 'Task: tinder_dog, titanic, simple_face, cifar10 or lus_covid', defaultValue: 'tinder_dog' }, numberOfUsers: { type: Number, alias: 'u', description: 'Number of users', defaultValue: 2 }, @@ -106,56 +80,6 @@ const unsafeArgs = parse( 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 }, -======= - testID: { - type: String, - alias: "i", - description: "ID of the testcase", - }, - task: { - type: String, - alias: "t", - description: - "Task: tinder_dog, titanic, simple_face, cifar10 or lus_covid", - defaultValue: "tinder_dog", - }, - numberOfUsers: { - type: Number, - alias: "u", - description: "Number of users", - defaultValue: 2, - }, - epochs: { - type: Number, - alias: "e", - description: "Number of epochs", - defaultValue: 10, - }, - roundDuration: { - type: Number, - alias: "r", - description: "Round duration (in epochs)", - defaultValue: 2, - }, - batchSize: { - type: Number, - alias: "b", - description: "Training batch size", - defaultValue: 10, - }, - validationSplit: { - type: Number, - alias: "v", - description: "Validation dataset ratio", - defaultValue: 0.2, - }, - save: { - type: Boolean, - alias: "s", - description: "Save logs of benchmark", - defaultValue: false, - }, ->>>>>>> develop host: { type: (raw: string) => new URL(raw), typeLabel: "URL", diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 4a9d5ebe3..28f4a3121 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -22,7 +22,6 @@ import { client as clients, } from "@epfml/discojs"; -<<<<<<< HEAD import { saveModelToDisk } from "@epfml/discojs-node"; import { getTaskData } from './data.js' import { args } from './args.js' @@ -87,27 +86,6 @@ async function runUser( }); const dir = getOutputDir(); -======= -import { getTaskData } from "./data.js"; -import { args } from "./args.js"; -import { makeUserLogFile } from "./user_log.js"; -import type { UserLogFile } from "./user_log.js"; - -async function runUser( - task: Task, - url: URL, - data: Dataset, - userIndex: number, - numberOfUsers: number, -): Promise> { - // cast as typescript isn't good with generics - const trainingScheme = task.trainingInformation.scheme as N; - const aggregator = aggregators.getAggregator(task); - const client = clients.getClient(trainingScheme, url, task, aggregator); - const disco = new Disco(task, client, { scheme: trainingScheme }); - - const dir = path.join(".", `${args.testID}`); ->>>>>>> develop await fs.mkdir(dir, { recursive: true }); const streamPath = path.join(dir, `client${userIndex}_local_log.jsonl`); @@ -115,7 +93,6 @@ async function runUser( // create a write stream that saves learning logs during the train let jsonStream: ReturnType | null = null; -<<<<<<< HEAD if (args.saveLogs){ jsonStream = createWriteStream(streamPath, {flags: "w"}); } @@ -126,14 +103,6 @@ async function runUser( let lastCheckpointRound: number | undefined = undefined; for await (const log of disco.trainSummary(data, validationData)){ -======= - if (args.save) { - jsonStream = createWriteStream(streamPath, { flags: "w" }); - } - - try { - for await (const log of disco.trainSummary(data)) { ->>>>>>> develop finalLog.push(log); if (jsonStream) { @@ -148,7 +117,6 @@ async function runUser( } debug(`Training took ${Date.now() - trainStart}ms for client ${userIndex}`); -<<<<<<< HEAD await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish // Save the trained model if requested if (args.saveModel) { @@ -158,27 +126,13 @@ async function runUser( runGarbageCollection(`client ${userIndex} final model save`); console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); } -======= - await new Promise((res, _) => setTimeout(() => res("timeout"), 1000)); // Wait for other peers to finish - ->>>>>>> develop // saving the entire per-user logs if (args.saveLogs) { const finalPath = path.join(dir, `client${userIndex}_local_log.json`); -<<<<<<< HEAD const clientId = trainingScheme === "local" ? `local-client-${userIndex}` : client.ownId; const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, clientId, finalLog); -======= - const userLog: UserLogFile = makeUserLogFile( - task, - numberOfUsers, - userIndex, - client.ownId, - finalLog, - ); ->>>>>>> develop await fs.writeFile(finalPath, JSON.stringify(userLog, null, 2)); } @@ -224,26 +178,8 @@ async function main( console.log({ args }); const dataSplits = await Promise.all( -<<<<<<< HEAD Range(0, numberOfUsers).map(async i => getTaskData(task.id, i, numberOfUsers, args.datasetPath)) ) -======= - Range(0, numberOfUsers).map(async (i) => - getTaskData(task.id, i, numberOfUsers), - ), - ); - const logs = await Promise.all( - dataSplits.map((data, i) => - runUser( - task, - args.host, - data as Dataset, - i, - numberOfUsers, - ), - ), - ); ->>>>>>> develop let validationData: Dataset | undefined = undefined; if (args.validationDatasetPath) { diff --git a/cli/src/data.ts b/cli/src/data.ts index 8bc43da84..f170460d4 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -1,7 +1,6 @@ import path from "node:path"; import { createReadStream } from "node:fs"; import { Dataset, processing } from "@epfml/discojs"; -<<<<<<< HEAD import { DataFormat, DataType, @@ -61,16 +60,6 @@ function loadTextSamples( } async function loadSimpleFaceData(userIdx: number, totalClient: number): Promise> { -======= -import { DataFormat, DataType, Image, Task } from "@epfml/discojs"; -import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node"; -import { Repeat } from "immutable"; - -async function loadSimpleFaceData( - userIdx: number, - totalClient: number, -): Promise> { ->>>>>>> develop const folder = path.join("..", "datasets", "simple_face"); const [adults, childs]: Dataset<[Image, string]>[] = [ @@ -159,18 +148,12 @@ function loadData( } export async function getTaskData( -<<<<<<< HEAD taskID: Task.ID, userIdx: number, totalClient: number, datasetPath?: string, isValidation?: boolean, validationDatasetPath?: string -======= - taskID: Task.ID, - userIdx: number, - totalClient: number, ->>>>>>> develop ): Promise> { switch (taskID) { case "simple_face": // remove diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index ffbb394fc..3fc5451cd 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -165,21 +165,12 @@ export abstract class Client extends EventEmitter<{ `[${shortenId(this.ownId)}] received EnoughParticipants message from server`, ); // Emit the last status emitted before waiting if defined -<<<<<<< HEAD if (this.#previousStatus !== undefined) this.emit("status", this.#previousStatus) this.nbOfParticipants = event.nbOfParticipants this.promiseForMoreParticipants = undefined resolve() }) }) -======= - if (this.#previousStatus !== undefined) - this.emit("status", this.#previousStatus); - this.nbOfParticipants = event.nbOfParticipants; - resolve(); - }); - }); ->>>>>>> develop } protected async waitForParticipantsIfNeeded(): Promise { diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index 36c1c7d58..fda02f5f2 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -121,19 +121,12 @@ export class WebSocketServer static async connect( url: URL, validateReceived: (msg: unknown) => msg is Message, -<<<<<<< HEAD validateSent: (msg: Message) => boolean): Promise { const ws = new WebSocket(url, { // Federated GPT updates can exceed the default ws payload limit. maxPayload: 1024 * 1024 * 1024, }) ws.binaryType = 'arraybuffer' -======= - validateSent: (msg: Message) => boolean, - ): Promise { - const ws = new WebSocket(url); - ws.binaryType = "arraybuffer"; ->>>>>>> develop const server: WebSocketServer = new WebSocketServer(ws, validateSent); @@ -158,20 +151,11 @@ export class WebSocketServer return await new Promise((resolve, reject) => { ws.onerror = (err: WebSocket.ErrorEvent) => { -<<<<<<< HEAD debug("websocket error while connecting/receiving: %o", err.message) reject(new Error(`Server unreachable: ${err.message}`)) } ws.onopen = () => { resolve(server) } }) -======= - reject(new Error(`Server unreachable: ${err.message}`)); - }; - ws.onopen = () => { - resolve(server); - }; - }); ->>>>>>> develop } disconnect(): Promise { diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 7b62c4217..ccd13d66b 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -93,21 +93,12 @@ export class FederatedClient extends Client<"federated"> { this.nbOfParticipants = nbOfParticipants; // Upon connecting, the server answers with a boolean // which indicates whether there are enough participants or not -<<<<<<< HEAD debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants) if (payload != null) { const latestWeights = serialization.weights.decode(payload) model.weights = latestWeights } return model -======= - debug( - `[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, - this.waitingForMoreParticipants, - ); - model.weights = serialization.weights.decode(payload); - return model; ->>>>>>> develop } /** @@ -154,7 +145,6 @@ export class FederatedClient extends Client<"federated"> { this.saveAndEmit("updating model"); // Send our local contribution to the server // and receive the server global update for this round as an answer to our contribution -<<<<<<< HEAD const payloadToServer = this.aggregator .makePayloads(weights) .get(SERVER_NODE_ID); @@ -194,34 +184,6 @@ export class FederatedClient extends Client<"federated"> { } = await waitMessage( this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update this.nbOfParticipants = nbOfParticipants // Save the current participants debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after receive`); -======= - const payloadToServer = this.aggregator - .makePayloads(weights) - .get(SERVER_NODE_ID); - if (payloadToServer === undefined) - throw new Error("aggregator didn't make a payload for the server"); - const msg: messages.SendPayload = { - type: type.SendPayload, - payload: await serialization.weights.encode(payloadToServer), - round: this.aggregator.round, - }; - - // Need to await the resulting global model right after sending our local contribution - // to make sure we don't miss it - debug( - `[${shortenId(this.ownId)}] sent its local update to the server for round ${this.aggregator.round}`, - ); - this.server.send(msg); - debug( - `[${shortenId(this.ownId)}] is waiting for server update for round ${this.aggregator.round + 1}`, - ); - const { - payload: payloadFromServer, - round: serverRound, - nbOfParticipants, - } = await waitMessage(this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update - this.nbOfParticipants = nbOfParticipants; // Save the current participants ->>>>>>> develop const serverResult = serialization.weights.decode(payloadFromServer); debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after decode server payload`); this.aggregator.setRound(serverRound); diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts index e22715d27..32ec504de 100644 --- a/discojs/src/client/federated/messages.ts +++ b/discojs/src/client/federated/messages.ts @@ -19,21 +19,12 @@ export type MessageFederated = | EnoughParticipants; export interface NewFederatedNodeInfo { -<<<<<<< HEAD type: type.NewFederatedNodeInfo id: NodeID waitForMoreParticipants: boolean payload?: serialization.Encoded | null; round: number nbOfParticipants: number -======= - type: type.NewFederatedNodeInfo; - id: NodeID; - waitForMoreParticipants: boolean; - payload: serialization.Encoded; - round: number; - nbOfParticipants: number; ->>>>>>> develop } export interface SendPayload { diff --git a/discojs/src/client/local_client.ts b/discojs/src/client/local_client.ts index bc0879be1..2d87292d4 100644 --- a/discojs/src/client/local_client.ts +++ b/discojs/src/client/local_client.ts @@ -9,17 +9,9 @@ export class LocalClient extends Client<"local"> { override onRoundBeginCommunication(): Promise { return Promise.resolve(); } -<<<<<<< HEAD // Return clones so the trainer can dispose the communication result without // disposing tensors owned by the model. override onRoundEndCommunication(weights: WeightsContainer): Promise { return Promise.resolve(weights.map((weight) => weight.clone())); -======= - // Simply return the local weights - override onRoundEndCommunication( - weights: WeightsContainer, - ): Promise { - return Promise.resolve(weights); ->>>>>>> develop } } diff --git a/discojs/src/default_tasks/index.ts b/discojs/src/default_tasks/index.ts index c4d389180..4d359574d 100644 --- a/discojs/src/default_tasks/index.ts +++ b/discojs/src/default_tasks/index.ts @@ -1,4 +1,3 @@ -<<<<<<< HEAD export { cifar10 } from './cifar10.js' export { lusCovid } from './lus_covid.js' export { mnist } from './mnist.js' @@ -8,12 +7,3 @@ export { wikitext } from './wikitext.js' export { tinderDog } from './tinder_dog.js' export { privacyrun } from './privacyrun.js' export { centralizedGPT2FineTune } from './centralized_gpt2_finetune.js' -======= -export { cifar10 } from "./cifar10.js"; -export { lusCovid } from "./lus_covid.js"; -export { mnist } from "./mnist.js"; -export { simpleFace } from "./simple_face.js"; -export { titanic } from "./titanic.js"; -export { wikitext } from "./wikitext.js"; -export { tinderDog } from "./tinder_dog.js"; ->>>>>>> develop diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index d60371138..c01b7aacb 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -8,7 +8,6 @@ type GPTModelType = | "gpt-nano"; export type GPTConfig = { -<<<<<<< HEAD lr: number contextLength: number vocabSize?: number @@ -35,27 +34,6 @@ export type GoldfishLossConfig = { h: number padTokenId?: number } -======= - lr: number; - contextLength: number; - vocabSize?: number; - modelType: GPTModelType; - evaluate?: boolean; - maxEvalBatches?: number; - evaluateEvery?: number; - maxIter?: number; - weightDecay?: number; - verbose?: 0 | 1; - debug?: boolean; - attnDrop?: number; - residDrop?: number; - embdDrop?: number; - nLayer?: number; - nHead?: number; - nEmbd?: number; - seed?: number; -}; ->>>>>>> develop // for a benchmark of performance, see https://github.com/epfml/disco/pull/659 export const DefaultGPTConfig: Required = { lr: 0.001, diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 9471150dd..8f32c2efb 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -14,13 +14,8 @@ import { BatchLogs, Model, EpochLogs } from "../index.js"; import { GPTModel } from "./model.js"; import evaluate from "./evaluate.js"; -<<<<<<< HEAD import { DefaultGPTConfig, DefaultGenerationConfig } from './config.js' import type { GoldfishLossConfig, GPTConfig, GenerationConfig } from './config.js' -======= -import { DefaultGPTConfig, DefaultGenerationConfig } from "./config.js"; -import type { GPTConfig, GenerationConfig } from "./config.js"; ->>>>>>> develop const debug = createDebug("discojs:models:gpt"); diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index a3743207e..70d6ef8aa 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -1,19 +1,11 @@ import createDebug from "debug"; import * as tf from "@tensorflow/tfjs"; -<<<<<<< HEAD import type { GoldfishLossConfig, GPTConfig } from './config.js' import { getModelSizes, DefaultGPTConfig } from './config.js' import { getCustomAdam, clipByGlobalNormObj } from './optimizers.js' import evaluate from './evaluate.js' import { GPTArchitecture } from './layers.js' -======= -import type { GPTConfig } from "./config.js"; -import { getModelSizes, DefaultGPTConfig } from "./config.js"; -import { getCustomAdam, clipByGlobalNormObj } from "./optimizers.js"; -import evaluate from "./evaluate.js"; -import { GPTArchitecture } from "./layers.js"; ->>>>>>> develop const debug = createDebug("discojs:models:gpt:model"); @@ -46,13 +38,9 @@ export declare abstract class Dataset { * */ export class GPTModel extends tf.LayersModel { -<<<<<<< HEAD protected readonly config: Required #debugLabel?: string #goldfishLoss?: GoldfishLossConfig -======= - protected readonly config: Required; ->>>>>>> develop constructor( partialConfig?: Partial, @@ -107,7 +95,6 @@ export class GPTModel extends tf.LayersModel { : tf.train.adam(this.config.lr); } -<<<<<<< HEAD setLearningRate(lr: number): void { this.config.lr = lr; this.optimizer?.dispose(); @@ -138,37 +125,8 @@ export class GPTModel extends tf.LayersModel { let weightUpdateTime = performance.now() await callbacks.onEpochBegin?.(epoch) const { xs, ys } = next.value as { xs: tf.Tensor2D, ys: tf.Tensor3D } -======= - override async fitDataset( - dataset: Dataset, - trainingArgs: tf.ModelFitDatasetArgs, - ): Promise { - const callbacks = trainingArgs.callbacks as tf.CustomCallbackArgs; - const evalDataset = trainingArgs.validationData as tf.data.Dataset<{ - xs: tf.Tensor2D; - ys: tf.Tensor3D; - }>; - await callbacks.onTrainBegin?.(); - - for (let epoch = 1; epoch <= trainingArgs.epochs; epoch++) { - let accuracyFraction: [number, number] = [0, 0]; - let averageLoss = 0; - let iteration = 1; - const iterator = await dataset.iterator(); - let next = await iterator.next(); - - while (next.done !== true && iteration <= this.config.maxIter) { - let weightUpdateTime = performance.now(); - await callbacks.onEpochBegin?.(epoch); - const { xs, ys } = next.value as { xs: tf.Tensor2D; ys: tf.Tensor3D }; - - let preprocessingTime = performance.now(); - await Promise.all([xs.data(), ys.data()]); - preprocessingTime = performance.now() - preprocessingTime; ->>>>>>> develop // TODO include as a tensor inside the model -<<<<<<< HEAD // const accTensor = tf.tidy(() => { // const logits = this.apply(xs) // if (Array.isArray(logits)) @@ -214,54 +172,10 @@ export class GPTModel extends tf.LayersModel { const loss = await lossTensor.array() averageLoss += loss weightUpdateTime = performance.now() - weightUpdateTime -======= - const accTensor = tf.tidy(() => { - const logits = this.apply(xs); - if (Array.isArray(logits)) - throw new Error("model outputs too many tensor"); - if (logits instanceof tf.SymbolicTensor) - throw new Error("model outputs symbolic tensor"); - return tf.metrics.categoricalAccuracy(ys, logits); - }); - const accSize = accTensor.shape.reduce((l, r) => l * r, 1); - const accSumTensor = accTensor.sum(); - const accSum = await accSumTensor.array(); - tf.dispose(accSumTensor); - if (typeof accSum !== "number") - throw new Error("got multiple accuracy sum"); - accuracyFraction = [ - accuracyFraction[0] + accSum, - accuracyFraction[1] + accSize, - ]; - tf.dispose([accTensor]); - - const lossTensor = tf.tidy(() => { - const { grads, value: lossTensor } = this.optimizer.computeGradients( - () => { - const logits = this.apply(xs); - if (Array.isArray(logits)) - throw new Error("model outputs too many tensor"); - if (logits instanceof tf.SymbolicTensor) - throw new Error("model outputs symbolic tensor"); - return tf.losses.softmaxCrossEntropy(ys, logits); - }, - ); - const gradsClipped = clipByGlobalNormObj(grads, 1); - this.optimizer.applyGradients(gradsClipped); - return lossTensor; - }); - - const loss = await lossTensor.array(); - averageLoss += loss; - weightUpdateTime = performance.now() - weightUpdateTime; - - tf.dispose([xs, ys, lossTensor]); ->>>>>>> develop if ( evalDataset !== undefined && this.config.evaluateEvery !== undefined && -<<<<<<< HEAD // iteration % this.config.evaluateEvery == 0 reportedIteration % this.config.evaluateEvery == 0 ){ @@ -270,19 +184,6 @@ export class GPTModel extends tf.LayersModel { } const memory = tf.memory().numBytes / 1024 / 1024 / 1024 debug(this.#debugMessage("training metrics: %O"), { -======= - iteration % this.config.evaluateEvery == 0 - ) { - const iterationLogs = await evaluate( - this, - evalDataset, - this.config.maxEvalBatches, - ); - debug("evaluation metrics: %O", iterationLogs); - } - const memory = tf.memory().numBytes / 1024 / 1024 / 1024; - debug("training metrics: %O", { ->>>>>>> develop epoch, iteration: reportedIteration, loss, diff --git a/discojs/src/models/index.ts b/discojs/src/models/index.ts index af7ec7e01..2f2119dc7 100644 --- a/discojs/src/models/index.ts +++ b/discojs/src/models/index.ts @@ -2,16 +2,10 @@ export { Model } from "./model.js"; export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js"; export { Tokenizer } from "./tokenizer.js"; -<<<<<<< HEAD export { GPT } from './gpt/index.js' export { ONNXModel } from './onnx.js' export { GPTConfig } from './gpt/config.js' export { GoldfishLossConfig } from './gpt/config.js' -======= -export { GPT } from "./gpt/index.js"; -export { ONNXModel } from "./onnx.js"; -export { GPTConfig } from "./gpt/config.js"; ->>>>>>> develop export { evaluate as evaluate_hellaswag, HellaSwagDataset, diff --git a/discojs/src/privacy.ts b/discojs/src/privacy.ts index 9b857872f..54804cd84 100644 --- a/discojs/src/privacy.ts +++ b/discojs/src/privacy.ts @@ -6,17 +6,11 @@ import type { WeightNormHistory } from "./training/trainer.js"; /** Computes the Frobenius norm of the given weights. */ export async function frobeniusNorm(weights: tf.Tensor): Promise { -<<<<<<< HEAD const squaredTensor = tf.tidy(() => weights.square().sum()); const squared = await squaredTensor.data(); squaredTensor.dispose(); if (squared.length !== 1) throw new Error("unexpected weights shape"); return Math.sqrt(squared[0]); -======= - const squared = await weights.square().sum().data(); - if (squared.length !== 1) throw new Error("unexpected weights shape"); - return Math.sqrt(squared[0]); ->>>>>>> develop } /** ALDP-FL implementation */ @@ -60,7 +54,6 @@ export async function addOptimalNoise( ); const clippedWeights = await clipNorm(weightUpdates, clippingRadius); -<<<<<<< HEAD try { return clippedWeights.map((w, i) => tf.tidy(() => w.add(tf.randomNormal(w.shape, 0, sigmas[i]))) @@ -68,11 +61,6 @@ export async function addOptimalNoise( } finally { clippedWeights.dispose(); } -======= - return clippedWeights.map((w, i) => - w.add(tf.randomNormal(w.shape, 0, sigmas[i])), - ); ->>>>>>> develop } /** From 6c53b6f287e6652c23085ce90396e5042cbb7ca2 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 8 Jul 2026 15:51:50 +0200 Subject: [PATCH 55/89] fix formatting --- cli/package.json | 8 +- cli/src/args.ts | 241 +++-- cli/src/cli.ts | 111 ++- cli/src/data.ts | 33 +- cli/src/evaluate_finetuned_gpt2.ts | 795 +++++++++-------- .../evaluate_finetuned_gpt2_full_answer.ts | 841 +++++++++--------- cli/src/measure_memorization_gpt2.ts | 148 ++- discojs/src/aggregator/mean.ts | 7 +- discojs/src/client/client.ts | 23 +- discojs/src/client/event_connection.ts | 28 +- .../src/client/federated/federated_client.ts | 59 +- discojs/src/client/federated/messages.ts | 10 +- discojs/src/client/local_client.ts | 4 +- .../centralized_gpt2_finetune.ts | 44 +- discojs/src/default_tasks/index.ts | 18 +- discojs/src/default_tasks/privacyrun.ts | 46 +- discojs/src/models/gpt/config.ts | 48 +- discojs/src/models/gpt/index.ts | 20 +- discojs/src/models/gpt/model.ts | 191 ++-- discojs/src/models/index.ts | 8 +- discojs/src/privacy.ts | 14 +- discojs/src/serialization/model.ts | 39 +- discojs/src/task/training_information.ts | 102 +-- discojs/src/training/disco.ts | 104 ++- discojs/src/training/trainer.ts | 193 ++-- onnx-converter/src/convert_onnx.ts | 19 +- .../src/controllers/federated_controller.ts | 135 +-- server/src/routes/training_router.ts | 12 +- server/tests/e2e/federated.spec.ts | 38 +- 29 files changed, 1896 insertions(+), 1443 deletions(-) diff --git a/cli/package.json b/cli/package.json index 3c8f9fd54..175d204f7 100644 --- a/cli/package.json +++ b/cli/package.json @@ -9,10 +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", + "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" }, diff --git a/cli/src/args.ts b/cli/src/args.ts index 11aa3982f..43ecda229 100644 --- a/cli/src/args.ts +++ b/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, models } from '@epfml/discojs' +import { defaultTasks, models } from "@epfml/discojs"; type AggregationStrategy = "mean" | "byzantine" | "secure"; @@ -12,23 +12,23 @@ function parseAggregator(raw: string): AggregationStrategy { } export interface BenchmarkArguments { - provider: TaskProvider; - testID: string - 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 + provider: TaskProvider; + testID: string; + 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; @@ -43,43 +43,136 @@ export interface BenchmarkArguments { // Secure aggregator maxShareValue?: number; - saveLogs: boolean - saveModel: boolean - saveCheckpoints: boolean - host: URL + saveLogs: boolean; + saveModel: boolean; + saveCheckpoints: boolean; + host: URL; } -type BenchmarkUnsafeArguments = Omit & { - task: string - datasetPath?: string - validationDatasetPath?: string - help?: boolean -} +type BenchmarkUnsafeArguments = Omit & { + 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( { - testID: { type: String, alias: 'i', description: 'ID of the testcase' }, - task: { type: String, alias: 't', description: 'Task: tinder_dog, titanic, simple_face, cifar10 or lus_covid', defaultValue: 'tinder_dog' }, - numberOfUsers: { type: Number, alias: 'u', description: 'Number of users', defaultValue: 2 }, - epochs: { type: Number, alias: 'e', description: 'Number of epochs', defaultValue: 10 }, - roundDuration: { type: Number, alias: 'r', 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', description: 'Training batch size', defaultValue: 10 }, - validationSplit : { type: Number, alias: 'v', description: 'Validation dataset ratio', defaultValue: 0.2 }, - 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 ./', 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 }, + testID: { type: String, alias: "i", description: "ID of the testcase" }, + task: { + type: String, + alias: "t", + description: + "Task: tinder_dog, titanic, simple_face, cifar10 or lus_covid", + defaultValue: "tinder_dog", + }, + numberOfUsers: { + type: Number, + alias: "u", + description: "Number of users", + defaultValue: 2, + }, + epochs: { + type: Number, + alias: "e", + description: "Number of epochs", + defaultValue: 10, + }, + roundDuration: { + type: Number, + alias: "r", + 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", + description: "Training batch size", + defaultValue: 10, + }, + validationSplit: { + type: Number, + alias: "v", + description: "Validation dataset ratio", + defaultValue: 0.2, + }, + 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 ./", + 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", @@ -194,17 +287,24 @@ export const args: BenchmarkArguments = { task.trainingInformation.roundDuration = unsafeArgs.roundDuration; task.trainingInformation.epochs = unsafeArgs.epochs; task.trainingInformation.validationSplit = unsafeArgs.validationSplit; - (task.trainingInformation as typeof task.trainingInformation & { - roundIterations?: number; - validationFrequency?: number; - }).roundIterations = unsafeArgs.roundIterations; - (task.trainingInformation as typeof task.trainingInformation & { - roundIterations?: number; - validationFrequency?: number; - }).validationFrequency = unsafeArgs.validationFrequency; + ( + task.trainingInformation as typeof task.trainingInformation & { + roundIterations?: number; + validationFrequency?: number; + } + ).roundIterations = unsafeArgs.roundIterations; + ( + task.trainingInformation as typeof task.trainingInformation & { + roundIterations?: number; + validationFrequency?: number; + } + ).validationFrequency = unsafeArgs.validationFrequency; if (unsafeArgs.goldfishLoss) { - if (task.dataType !== "text" || task.trainingInformation.tensorBackend !== "gpt") + 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"); @@ -220,9 +320,17 @@ export const args: BenchmarkArguments = { } 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) + 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; @@ -305,12 +413,19 @@ export const args: BenchmarkArguments = { 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 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}`); + console.log( + `Overriding GPT learning rate to ${unsafeArgs.learningRate}`, + ); } return model; diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 28f4a3121..6a6b40aa9 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -23,8 +23,8 @@ import { } from "@epfml/discojs"; import { saveModelToDisk } from "@epfml/discojs-node"; -import { getTaskData } from './data.js' -import { args } from './args.js' +import { getTaskData } from "./data.js"; +import { args } from "./args.js"; import { makeUserLogFile } from "./user_log.js"; import type { UserLogFile } from "./user_log.js"; @@ -39,7 +39,10 @@ function getOutputDir(): string { function runGarbageCollection(label: string): void { const gc = (globalThis as typeof globalThis & { gc?: () => void }).gc; if (gc === undefined) { - debug("%s skipped explicit GC because node was not started with --expose-gc", label); + debug( + "%s skipped explicit GC because node was not started with --expose-gc", + label, + ); return; } gc(); @@ -50,11 +53,17 @@ async function saveClientModelCheckpoint( userIndex: number, round: number, ): Promise { - const checkpointDir = path.join(getOutputDir(), "checkpoints", `round_${round}`); + const checkpointDir = path.join( + getOutputDir(), + "checkpoints", + `round_${round}`, + ); const checkpointFileName = `client${userIndex}_model.json`; await saveModelToDisk(model, checkpointDir, checkpointFileName); - console.log(`Checkpoint saved for client ${userIndex} round ${round} at ${checkpointDir}/${checkpointFileName}`); + console.log( + `Checkpoint saved for client ${userIndex} round ${round} at ${checkpointDir}/${checkpointFileName}`, + ); } async function enqueueClientModelCheckpoint( @@ -62,29 +71,31 @@ async function enqueueClientModelCheckpoint( userIndex: number, round: number, ): Promise { - const save = checkpointQueue.then(() => saveClientModelCheckpoint(model, userIndex, round)); + const save = checkpointQueue.then(() => + saveClientModelCheckpoint(model, userIndex, round), + ); checkpointQueue = save.catch(() => undefined); await save; } async function runUser( - task: Task, - url: URL, - data: Dataset, + task: Task, + url: URL, + data: Dataset, validationData: Dataset | undefined, userIndex: number, numberOfUsers: number, ): Promise> { debug(`Starting runUser for client ${userIndex}`); - const trainingScheme = task.trainingInformation.scheme as N - const aggregator = aggregators.getAggregator(task) - const client = clients.getClient(trainingScheme, url, task, aggregator) + const trainingScheme = task.trainingInformation.scheme as N; + const aggregator = aggregators.getAggregator(task); + const client = clients.getClient(trainingScheme, url, task, aggregator); const disco = new Disco(task, client, { scheme: trainingScheme, preprocessOnce: false, debugLabel: `client${userIndex}`, }); - + const dir = getOutputDir(); await fs.mkdir(dir, { recursive: true }); const streamPath = path.join(dir, `client${userIndex}_local_log.jsonl`); @@ -93,16 +104,16 @@ async function runUser( // create a write stream that saves learning logs during the train let jsonStream: ReturnType | null = null; - if (args.saveLogs){ - jsonStream = createWriteStream(streamPath, {flags: "w"}); + if (args.saveLogs) { + jsonStream = createWriteStream(streamPath, { flags: "w" }); } - try{ + try { debug(`Starting training for client ${userIndex}`); const trainStart = Date.now(); let lastCheckpointRound: number | undefined = undefined; - for await (const log of disco.trainSummary(data, validationData)){ + for await (const log of disco.trainSummary(data, validationData)) { finalLog.push(log); if (jsonStream) { @@ -110,29 +121,43 @@ async function runUser( } if (args.saveCheckpoints && lastCheckpointRound !== log.round) { - await enqueueClientModelCheckpoint(disco.trainer.model, userIndex, log.round); - runGarbageCollection(`client ${userIndex} round ${log.round} checkpoint`); + await enqueueClientModelCheckpoint( + disco.trainer.model, + userIndex, + log.round, + ); + runGarbageCollection( + `client ${userIndex} round ${log.round} checkpoint`, + ); lastCheckpointRound = log.round; } } debug(`Training took ${Date.now() - trainStart}ms for client ${userIndex}`); - await new Promise((res, _) => setTimeout(() => res('timeout'), 1000)) // Wait for other peers to finish - // Save the trained model if requested - if (args.saveModel) { - const modelDir = path.join(getOutputDir(), "models"); - const modelFileName = `client${userIndex}_model.json`; - await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); - runGarbageCollection(`client ${userIndex} final model save`); - console.log(`Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`); - } + await new Promise((res, _) => setTimeout(() => res("timeout"), 1000)); // Wait for other peers to finish + // Save the trained model if requested + if (args.saveModel) { + const modelDir = path.join(getOutputDir(), "models"); + const modelFileName = `client${userIndex}_model.json`; + await saveModelToDisk(disco.trainer.model, modelDir, modelFileName); + runGarbageCollection(`client ${userIndex} final model save`); + console.log( + `Model saved for client ${userIndex} at ${modelDir}/${modelFileName}`, + ); + } // saving the entire per-user logs if (args.saveLogs) { const finalPath = path.join(dir, `client${userIndex}_local_log.json`); const clientId = trainingScheme === "local" ? `local-client-${userIndex}` : client.ownId; - const userLog: UserLogFile = makeUserLogFile(task, numberOfUsers, userIndex, clientId, finalLog); + const userLog: UserLogFile = makeUserLogFile( + task, + numberOfUsers, + userIndex, + clientId, + finalLog, + ); await fs.writeFile(finalPath, JSON.stringify(userLog, null, 2)); } @@ -178,19 +203,37 @@ async function main( console.log({ args }); const dataSplits = await Promise.all( - Range(0, numberOfUsers).map(async i => getTaskData(task.id, i, numberOfUsers, args.datasetPath)) - ) + Range(0, numberOfUsers).map(async (i) => + getTaskData(task.id, i, numberOfUsers, args.datasetPath), + ), + ); let validationData: Dataset | undefined = undefined; if (args.validationDatasetPath) { validationData = ( - await getTaskData(task.id, 0, 1, args.validationDatasetPath, true, args.validationDatasetPath) + await getTaskData( + task.id, + 0, + 1, + args.validationDatasetPath, + true, + args.validationDatasetPath, + ) ).cached() as Dataset; } const logs = await Promise.all( - dataSplits.map((data, i) => runUser(task, args.host, data as Dataset, validationData, i, numberOfUsers)) - ) + dataSplits.map((data, i) => + runUser( + task, + args.host, + data as Dataset, + validationData, + i, + numberOfUsers, + ), + ), + ); if (args.saveLogs) { const dir = path.join(getOutputDir(), `${task.id}`); diff --git a/cli/src/data.ts b/cli/src/data.ts index f170460d4..1bd68e0d7 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -1,13 +1,7 @@ import path from "node:path"; import { createReadStream } from "node:fs"; import { Dataset, processing } from "@epfml/discojs"; -import { - DataFormat, - DataType, - Image, - Task, - Text, -} from "@epfml/discojs"; +import { DataFormat, DataType, Image, Task, Text } from "@epfml/discojs"; import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node"; import { Repeat } from "immutable"; @@ -31,7 +25,9 @@ function loadTextSamples( let delimiterIndex = buffer.indexOf(sampleDelimiter); while (delimiterIndex !== -1) { - const sample = buffer.slice(0, delimiterIndex + sampleDelimiter.length).trim(); + const sample = buffer + .slice(0, delimiterIndex + sampleDelimiter.length) + .trim(); const shouldYield = userIdx === undefined || totalClient === undefined || @@ -59,7 +55,10 @@ function loadTextSamples( }); } -async function loadSimpleFaceData(userIdx: number, totalClient: number): Promise> { +async function loadSimpleFaceData( + userIdx: number, + totalClient: number, +): Promise> { const folder = path.join("..", "datasets", "simple_face"); const [adults, childs]: Dataset<[Image, string]>[] = [ @@ -148,12 +147,12 @@ function loadData( } export async function getTaskData( - taskID: Task.ID, - userIdx: number, + taskID: Task.ID, + userIdx: number, totalClient: number, datasetPath?: string, isValidation?: boolean, - validationDatasetPath?: string + validationDatasetPath?: string, ): Promise> { switch (taskID) { case "simple_face": // remove @@ -186,18 +185,16 @@ export async function getTaskData( const filePath = isValidation && validationDatasetPath ? validationDatasetPath - : datasetPath ?? "../datasets/med_mcq/train.txt"; + : (datasetPath ?? "../datasets/med_mcq/train.txt"); // Keep validation shared, but shard training data across clients by MCQ sample. if (isValidation) { return loadTextSamples(filePath) as Dataset; } - return loadTextSamples( - filePath, - userIdx, - totalClient, - ) as Dataset; + return loadTextSamples(filePath, userIdx, totalClient) as Dataset< + DataFormat.Raw[D] + >; } default: throw new Error(`Data loader for ${taskID} not implemented.`); diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index f6f01874c..4f82e2102 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -6,469 +6,488 @@ import { models, Tokenizer } from "@epfml/discojs"; import { loadModelFromDisk } from "@epfml/discojs-node"; interface Args { - modelPath: string; - testPath: string; - maxSamples?: number; - savePath?: string; - compareFormats?: boolean; - promptFormat?: PromptFormatName; - contextLength?: number; - help?: boolean; + modelPath: string; + testPath: string; + maxSamples?: number; + savePath?: string; + compareFormats?: boolean; + promptFormat?: PromptFormatName; + contextLength?: number; + help?: boolean; } // HOW TO RUN // npm -w cli run eval_finetuned_gpt2 -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/train_no_exp.txt --maxSamples 100 const PromptFormatNames = [ - "answer-colon-space", - "answer-colon", - "answer-newline", + "answer-colon-space", + "answer-colon", + "answer-newline", ] as const; -type PromptFormatName = typeof PromptFormatNames[number]; +type PromptFormatName = (typeof PromptFormatNames)[number]; type PromptFormat = { - name: PromptFormatName; - makePrompt: (basePrompt: string) => string; - makeContinuation: (option: string) => string; + name: PromptFormatName; + makePrompt: (basePrompt: string) => string; + makeContinuation: (option: string) => string; }; const promptFormats: PromptFormat[] = [ - { - name: "answer-colon-space", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, - makeContinuation: (option) => option, - }, - { - name: "answer-colon", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, - makeContinuation: (option) => ` ${option}`, - }, - { - name: "answer-newline", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, - makeContinuation: (option) => option, - }, + { + name: "answer-colon-space", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, + makeContinuation: (option) => option, + }, + { + name: "answer-colon", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, + makeContinuation: (option) => ` ${option}`, + }, + { + name: "answer-newline", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, + makeContinuation: (option) => option, + }, ]; function castPromptFormatName(raw: string): PromptFormatName { - for (const name of PromptFormatNames) { - if (raw === name) return name; - } - throw new Error(`Invalid promptFormat: ${raw}`); + for (const name of PromptFormatNames) { + if (raw === name) return name; + } + throw new Error(`Invalid promptFormat: ${raw}`); } function commonPrefixLength(left: number[], right: number[]): number { - const maxLength = Math.min(left.length, right.length); + const maxLength = Math.min(left.length, right.length); - for (let i = 0; i < maxLength; i++) { - if (left[i] !== right[i]) return i; - } + for (let i = 0; i < maxLength; i++) { + if (left[i] !== right[i]) return i; + } - return maxLength; + return maxLength; } -function predictTokenLogits(tfModel: tf.LayersModel, inputTensor: tf.Tensor2D): tf.Tensor3D { - const logits = tfModel.predict(inputTensor); - if (Array.isArray(logits)) { - throw new Error("Expected GPT model to return a single logits tensor"); - } - if (logits.rank !== 3) { - logits.dispose(); - throw new Error(`Expected GPT logits to have rank 3, got rank ${logits.rank}`); - } - return logits as tf.Tensor3D; +function predictTokenLogits( + tfModel: tf.LayersModel, + inputTensor: tf.Tensor2D, +): tf.Tensor3D { + const logits = tfModel.predict(inputTensor); + if (Array.isArray(logits)) { + throw new Error("Expected GPT model to return a single logits tensor"); + } + if (logits.rank !== 3) { + logits.dispose(); + throw new Error( + `Expected GPT logits to have rank 3, got rank ${logits.rank}`, + ); + } + return logits as tf.Tensor3D; } async function loadDataset(filePath: string, limit = -1): Promise { - const text = await fs.readFile(filePath, "utf-8"); - const samples = text - .split("<|endoftext|>") - .map((sample) => - sample - .replaceAll("<|startoftext|>", "") - .trim(), - ) - .filter((sample) => sample !== ""); - - return limit === -1 ? samples : samples.slice(0, limit); -} + const text = await fs.readFile(filePath, "utf-8"); + const samples = text + .split("<|endoftext|>") + .map((sample) => sample.replaceAll("<|startoftext|>", "").trim()) + .filter((sample) => sample !== ""); -function parseSample(sample: string) { - const lines = sample.split("\n"); - - let answer = ""; - const promptLines: string[] = []; - - for (const line of lines) { - const trimmed = line.trim(); - if (trimmed.startsWith("Answer:")) { - answer = trimmed.replace("Answer:", "").trim().charAt(0).toUpperCase(); - } else { - promptLines.push(line); - } - } - - const basePrompt = promptLines.join("\n").trim(); - return { basePrompt, answer }; + return limit === -1 ? samples : samples.slice(0, limit); } -async function scoreContinuations( - tfModel: tf.LayersModel, - tokenizer: Tokenizer, - prompt: string, - continuations: string[], - contextLength: number -): Promise<{ score: number; promptTokens: number; continuationTokens: number; usedInputTokens: number }[]> { - const promptTokens = tokenizer.tokenize(prompt).toArray(); - const scoredInputs = continuations.map((continuation) => { - const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); - const continuationStart = commonPrefixLength(promptTokens, fullTokens); - const continuationTokens = fullTokens.length - continuationStart; - const inputTokens = fullTokens.slice(0, -1); - const offset = Math.max(0, inputTokens.length - contextLength); - const truncatedInputTokens = inputTokens.slice(offset); - - return { - fullTokens, - continuationStart, - continuationTokens, - offset, - truncatedInputTokens, - }; - }); +function parseSample(sample: string) { + const lines = sample.split("\n"); - const maxInputLength = scoredInputs.reduce( - (maxLength, scoredInput) => Math.max(maxLength, scoredInput.truncatedInputTokens.length), - 0, - ); + let answer = ""; + const promptLines: string[] = []; - if (maxInputLength === 0) { - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); + for (const line of lines) { + const trimmed = line.trim(); + if (trimmed.startsWith("Answer:")) { + answer = trimmed.replace("Answer:", "").trim().charAt(0).toUpperCase(); + } else { + promptLines.push(line); } + } - const canScoreFromPromptOnly = scoredInputs.every( - ({ continuationStart, continuationTokens }) => - continuationStart === promptTokens.length && continuationTokens === 1, - ); + const basePrompt = promptLines.join("\n").trim(); + return { basePrompt, answer }; +} - if (canScoreFromPromptOnly) { - const offset = Math.max(0, promptTokens.length - contextLength); - const truncatedPromptTokens = promptTokens.slice(offset); - - if (truncatedPromptTokens.length === 0) { - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: truncatedPromptTokens.length, - })); - } - - const inputTensor = tf.tensor2d( - [truncatedPromptTokens], - [1, truncatedPromptTokens.length], - "int32", - ); - - const optionScores = tf.tidy(() => { - const logits = predictTokenLogits(tfModel, inputTensor); - const lastLogits = logits - .slice([0, truncatedPromptTokens.length - 1, 0], [1, 1, -1]) - .reshape([-1]); - const logProbs = tf.logSoftmax(lastLogits); - const continuationTokenIds = scoredInputs.map( - ({ fullTokens, continuationStart }) => fullTokens[continuationStart], - ); - return tf.gather(logProbs, continuationTokenIds); - }); - - const scores = await optionScores.array(); - - inputTensor.dispose(); - optionScores.dispose(); - - return scoredInputs.map((scoredInput, index) => ({ - score: scores[index], - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: truncatedPromptTokens.length, - })); +async function scoreContinuations( + tfModel: tf.LayersModel, + tokenizer: Tokenizer, + prompt: string, + continuations: string[], + contextLength: number, +): Promise< + { + score: number; + promptTokens: number; + continuationTokens: number; + usedInputTokens: number; + }[] +> { + const promptTokens = tokenizer.tokenize(prompt).toArray(); + const scoredInputs = continuations.map((continuation) => { + const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); + const continuationStart = commonPrefixLength(promptTokens, fullTokens); + const continuationTokens = fullTokens.length - continuationStart; + const inputTokens = fullTokens.slice(0, -1); + const offset = Math.max(0, inputTokens.length - contextLength); + const truncatedInputTokens = inputTokens.slice(offset); + + return { + fullTokens, + continuationStart, + continuationTokens, + offset, + truncatedInputTokens, + }; + }); + + const maxInputLength = scoredInputs.reduce( + (maxLength, scoredInput) => + Math.max(maxLength, scoredInput.truncatedInputTokens.length), + 0, + ); + + if (maxInputLength === 0) { + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + } + + const canScoreFromPromptOnly = scoredInputs.every( + ({ continuationStart, continuationTokens }) => + continuationStart === promptTokens.length && continuationTokens === 1, + ); + + if (canScoreFromPromptOnly) { + const offset = Math.max(0, promptTokens.length - contextLength); + const truncatedPromptTokens = promptTokens.slice(offset); + + if (truncatedPromptTokens.length === 0) { + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: truncatedPromptTokens.length, + })); } - const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ - ...truncatedInputTokens, - ...Array(maxInputLength - truncatedInputTokens.length).fill(0), - ]); - const inputTensor = tf.tensor2d( - paddedInputs, - [paddedInputs.length, maxInputLength], - "int32", + [truncatedPromptTokens], + [1, truncatedPromptTokens.length], + "int32", ); - const targetIndexes: number[][] = []; - const targetTokenIds: number[] = []; - const targetOwners: number[] = []; - - scoredInputs.forEach((scoredInput, batchIdx) => { - const { - fullTokens, - continuationStart, - offset, - truncatedInputTokens, - } = scoredInput; - - // Usually A/B/C/D is one token and the prompt-only fast path above handles it. - // Keep this fallback for prompt/continuation tokenizer merges and multi-token labels. - for (let targetPos = continuationStart; targetPos < fullTokens.length; targetPos++) { - const targetToken = fullTokens[targetPos]; - const logitPos = targetPos - 1 - offset; - if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; - targetIndexes.push([batchIdx, logitPos]); - targetTokenIds.push(targetToken); - targetOwners.push(batchIdx); - } + const optionScores = tf.tidy(() => { + const logits = predictTokenLogits(tfModel, inputTensor); + const lastLogits = logits + .slice([0, truncatedPromptTokens.length - 1, 0], [1, 1, -1]) + .reshape([-1]); + const logProbs = tf.logSoftmax(lastLogits); + const continuationTokenIds = scoredInputs.map( + ({ fullTokens, continuationStart }) => fullTokens[continuationStart], + ); + return tf.gather(logProbs, continuationTokenIds); }); - if (targetIndexes.length === 0) { - inputTensor.dispose(); - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); - } - - const logits = predictTokenLogits(tfModel, inputTensor); - const targetLogProbs = tf.tidy(() => { - const targetIndexTensor = tf.tensor2d( - targetIndexes, - [targetIndexes.length, 2], - "int32", - ); - const targetTokenIndexTensor = tf.tensor2d( - targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), - [targetTokenIds.length, 2], - "int32", - ); - const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; - const logProbs = tf.logSoftmax(targetLogits, -1); - return tf.gatherND(logProbs, targetTokenIndexTensor); - }); - - const targetScores = await targetLogProbs.array() as number[]; - const scoreSums = Array(scoredInputs.length).fill(0) as number[]; - const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; - - targetScores.forEach((score, index) => { - const owner = targetOwners[index]; - scoreSums[owner] += score; - scoreCounts[owner]++; - }); + const scores = await optionScores.array(); - const results = scoredInputs.map((scoredInput, index) => { - return { - score: scoreCounts[index] === 0 - ? Number.NEGATIVE_INFINITY - : scoreSums[index] / scoreCounts[index], - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - }; - }); + inputTensor.dispose(); + optionScores.dispose(); + + return scoredInputs.map((scoredInput, index) => ({ + score: scores[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: truncatedPromptTokens.length, + })); + } + + const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ + ...truncatedInputTokens, + ...Array(maxInputLength - truncatedInputTokens.length).fill(0), + ]); + + const inputTensor = tf.tensor2d( + paddedInputs, + [paddedInputs.length, maxInputLength], + "int32", + ); + + const targetIndexes: number[][] = []; + const targetTokenIds: number[] = []; + const targetOwners: number[] = []; + + scoredInputs.forEach((scoredInput, batchIdx) => { + const { fullTokens, continuationStart, offset, truncatedInputTokens } = + scoredInput; + + // Usually A/B/C/D is one token and the prompt-only fast path above handles it. + // Keep this fallback for prompt/continuation tokenizer merges and multi-token labels. + for ( + let targetPos = continuationStart; + targetPos < fullTokens.length; + targetPos++ + ) { + const targetToken = fullTokens[targetPos]; + const logitPos = targetPos - 1 - offset; + if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; + targetIndexes.push([batchIdx, logitPos]); + targetTokenIds.push(targetToken); + targetOwners.push(batchIdx); + } + }); + if (targetIndexes.length === 0) { inputTensor.dispose(); - logits.dispose(); - targetLogProbs.dispose(); + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + } + + const logits = predictTokenLogits(tfModel, inputTensor); + const targetLogProbs = tf.tidy(() => { + const targetIndexTensor = tf.tensor2d( + targetIndexes, + [targetIndexes.length, 2], + "int32", + ); + const targetTokenIndexTensor = tf.tensor2d( + targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), + [targetTokenIds.length, 2], + "int32", + ); + const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; + const logProbs = tf.logSoftmax(targetLogits, -1); + return tf.gatherND(logProbs, targetTokenIndexTensor); + }); + + const targetScores = (await targetLogProbs.array()) as number[]; + const scoreSums = Array(scoredInputs.length).fill(0) as number[]; + const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; + + targetScores.forEach((score, index) => { + const owner = targetOwners[index]; + scoreSums[owner] += score; + scoreCounts[owner]++; + }); + + const results = scoredInputs.map((scoredInput, index) => { + return { + score: + scoreCounts[index] === 0 + ? Number.NEGATIVE_INFINITY + : scoreSums[index] / scoreCounts[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + }; + }); - return results; + inputTensor.dispose(); + logits.dispose(); + targetLogProbs.dispose(); + + return results; } async function benchmarkQA( - model: models.GPT, - tokenizer: Tokenizer, - dataset: string[], - format: PromptFormat, - contextLength: number, - savePath?: string + model: models.GPT, + tokenizer: Tokenizer, + dataset: string[], + format: PromptFormat, + contextLength: number, + savePath?: string, ): Promise { - console.log(`=== QA LOGPROB BENCHMARK (${format.name}) ===`); - console.log(`Context length: ${contextLength}`); + console.log(`=== QA LOGPROB BENCHMARK (${format.name}) ===`); + console.log(`Context length: ${contextLength}`); - const tfModel = model.extract(); + const tfModel = model.extract(); - let correct = 0; - let total = 0; + let correct = 0; + let total = 0; - const options = ["A", "B", "C", "D"]; + const options = ["A", "B", "C", "D"]; - const confusion: Record> = { - A: { A: 0, B: 0, C: 0, D: 0 }, - B: { A: 0, B: 0, C: 0, D: 0 }, - C: { A: 0, B: 0, C: 0, D: 0 }, - D: { A: 0, B: 0, C: 0, D: 0 } - }; + const confusion: Record> = { + A: { A: 0, B: 0, C: 0, D: 0 }, + B: { A: 0, B: 0, C: 0, D: 0 }, + C: { A: 0, B: 0, C: 0, D: 0 }, + D: { A: 0, B: 0, C: 0, D: 0 }, + }; - type PredictionLog = { - predicted: string; - answer: string; - correct: boolean; - scores: Record; - promptTokens: number; - continuationTokens: Record; - usedInputTokens: Record; - }; + type PredictionLog = { + predicted: string; + answer: string; + correct: boolean; + scores: Record; + promptTokens: number; + continuationTokens: Record; + usedInputTokens: Record; + }; - const logs: PredictionLog[] = []; - - const start = Date.now(); - - for (const sample of dataset) { - const { basePrompt, answer } = parseSample(sample); - - if (!options.includes(answer)) { - console.log("Invalid answer:", JSON.stringify(answer)); - continue; - } - - const prompt = format.makePrompt(basePrompt); - - const scores: number[] = []; - const continuationTokens: number[] = []; - const usedInputTokens: number[] = []; - const results = await scoreContinuations( - tfModel, - tokenizer, - prompt, - options.map((opt) => format.makeContinuation(opt)), - contextLength, - ); - - for (const result of results) { - scores.push(result.score); - continuationTokens.push(result.continuationTokens); - usedInputTokens.push(result.usedInputTokens); - } - - let bestIdx = 0; - for (let i = 1; i < scores.length; i++) { - if (scores[i] > scores[bestIdx]) bestIdx = i; - } - - const predicted = options[bestIdx]; - - if (predicted === answer) correct++; - total++; - - if (confusion[answer]?.[predicted] === undefined) { - throw new Error(`Unexpected confusion matrix key: answer=${answer}, predicted=${predicted}`); - } - confusion[answer][predicted]++; - - const scoreMap = Object.fromEntries( - options.map((opt, i) => [opt, scores[i]]) - ); - - logs.push({ - predicted, - answer, - correct: predicted === answer, - scores: scoreMap, - promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, - continuationTokens: Object.fromEntries( - options.map((opt, i) => [opt, continuationTokens[i]]) - ), - usedInputTokens: Object.fromEntries( - options.map((opt, i) => [opt, usedInputTokens[i]]) - ), - }); - - if (total % 50 === 0) { - console.log(`Processed ${total} samples...`); - } - } + const logs: PredictionLog[] = []; - const accuracy = correct / total; - const duration = ((Date.now() - start) / 1000).toFixed(2); + const start = Date.now(); - console.log("\n========================="); - console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); - console.log(`Time: ${duration}s`); - console.log("=========================\n"); + for (const sample of dataset) { + const { basePrompt, answer } = parseSample(sample); - console.log("Confusion Matrix:"); - console.table(confusion); + if (!options.includes(answer)) { + console.log("Invalid answer:", JSON.stringify(answer)); + continue; + } - console.log("\nPer-class accuracy:"); - for (const cls of options) { - const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); - const correctCls = confusion[cls][cls]; - const acc = totalCls ? (correctCls / totalCls) * 100 : 0; + const prompt = format.makePrompt(basePrompt); + + const scores: number[] = []; + const continuationTokens: number[] = []; + const usedInputTokens: number[] = []; + const results = await scoreContinuations( + tfModel, + tokenizer, + prompt, + options.map((opt) => format.makeContinuation(opt)), + contextLength, + ); - console.log(`${cls}: ${acc.toFixed(2)}%`); + for (const result of results) { + scores.push(result.score); + continuationTokens.push(result.continuationTokens); + usedInputTokens.push(result.usedInputTokens); } - if (savePath) { - await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); - console.log(`Saved results to ${savePath}`); + let bestIdx = 0; + for (let i = 1; i < scores.length; i++) { + if (scores[i] > scores[bestIdx]) bestIdx = i; } - return accuracy; -} + const predicted = options[bestIdx]; -async function main() { - const args = parse({ - modelPath: { type: String }, - testPath: { type: String }, - maxSamples: { type: Number, optional: true, defaultValue: 100 }, - savePath: { type: String, optional: true }, - compareFormats: { type: Boolean, optional: true, defaultValue: false }, - promptFormat: { - type: (raw: string) => castPromptFormatName(raw), - optional: true, - defaultValue: "answer-colon-space", - }, - contextLength: { type: Number, optional: true }, - help: { type: Boolean, optional: true } - }); + if (predicted === answer) correct++; + total++; - console.log("Loading tokenizer..."); - const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + if (confusion[answer]?.[predicted] === undefined) { + throw new Error( + `Unexpected confusion matrix key: answer=${answer}, predicted=${predicted}`, + ); + } + confusion[answer][predicted]++; - console.log("Loading model..."); - const model = await loadModelFromDisk(args.modelPath); + const scoreMap = Object.fromEntries( + options.map((opt, i) => [opt, scores[i]]), + ); - if (!(model instanceof models.GPT)) { - throw new Error("Model must be GPT"); + logs.push({ + predicted, + answer, + correct: predicted === answer, + scores: scoreMap, + promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, + continuationTokens: Object.fromEntries( + options.map((opt, i) => [opt, continuationTokens[i]]), + ), + usedInputTokens: Object.fromEntries( + options.map((opt, i) => [opt, usedInputTokens[i]]), + ), + }); + + if (total % 50 === 0) { + console.log(`Processed ${total} samples...`); } + } - console.log("Loading dataset..."); - const dataset = await loadDataset(args.testPath, args.maxSamples); + const accuracy = correct / total; + const duration = ((Date.now() - start) / 1000).toFixed(2); - console.log(`Loaded ${dataset.length} samples`); + console.log("\n========================="); + console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); + console.log(`Time: ${duration}s`); + console.log("=========================\n"); - const contextLength = args.contextLength ?? model.config.contextLength; - const formats = args.compareFormats - ? promptFormats - : promptFormats.filter((format) => format.name === args.promptFormat); + console.log("Confusion Matrix:"); + console.table(confusion); - for (const format of formats) { - const savePath = - args.savePath === undefined || formats.length === 1 - ? args.savePath - : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + console.log("\nPer-class accuracy:"); + for (const cls of options) { + const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); + const correctCls = confusion[cls][cls]; + const acc = totalCls ? (correctCls / totalCls) * 100 : 0; - await benchmarkQA(model, tokenizer, dataset, format, contextLength, savePath); - } + console.log(`${cls}: ${acc.toFixed(2)}%`); + } + + if (savePath) { + await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); + console.log(`Saved results to ${savePath}`); + } + + return accuracy; +} + +async function main() { + const args = parse({ + modelPath: { type: String }, + testPath: { type: String }, + maxSamples: { type: Number, optional: true, defaultValue: 100 }, + savePath: { type: String, optional: true }, + compareFormats: { type: Boolean, optional: true, defaultValue: false }, + promptFormat: { + type: (raw: string) => castPromptFormatName(raw), + optional: true, + defaultValue: "answer-colon-space", + }, + contextLength: { type: Number, optional: true }, + help: { type: Boolean, optional: true }, + }); + + console.log("Loading tokenizer..."); + const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + + console.log("Loading model..."); + const model = await loadModelFromDisk(args.modelPath); + + if (!(model instanceof models.GPT)) { + throw new Error("Model must be GPT"); + } + + console.log("Loading dataset..."); + const dataset = await loadDataset(args.testPath, args.maxSamples); + + console.log(`Loaded ${dataset.length} samples`); + + const contextLength = args.contextLength ?? model.config.contextLength; + const formats = args.compareFormats + ? promptFormats + : promptFormats.filter((format) => format.name === args.promptFormat); + + for (const format of formats) { + const savePath = + args.savePath === undefined || formats.length === 1 + ? args.savePath + : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + + await benchmarkQA( + model, + tokenizer, + dataset, + format, + contextLength, + savePath, + ); + } - console.log("Done."); + console.log("Done."); } main().catch(console.error); diff --git a/cli/src/evaluate_finetuned_gpt2_full_answer.ts b/cli/src/evaluate_finetuned_gpt2_full_answer.ts index 661ae9af2..5de9d827d 100644 --- a/cli/src/evaluate_finetuned_gpt2_full_answer.ts +++ b/cli/src/evaluate_finetuned_gpt2_full_answer.ts @@ -6,495 +6,526 @@ import { models, Tokenizer } from "@epfml/discojs"; import { loadModelFromDisk } from "@epfml/discojs-node"; interface Args { - modelPath: string; - testPath: string; - maxSamples?: number; - savePath?: string; - compareFormats?: boolean; - promptFormat?: PromptFormatName; - contextLength?: number; - help?: boolean; + modelPath: string; + testPath: string; + maxSamples?: number; + savePath?: string; + compareFormats?: boolean; + promptFormat?: PromptFormatName; + contextLength?: number; + help?: boolean; } // HOW TO RUN // npm -w cli run eval_finetuned_gpt2_full_answer -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/test.txt --maxSamples 100 const PromptFormatNames = [ - "answer-colon-space", - "answer-colon", - "answer-newline", + "answer-colon-space", + "answer-colon", + "answer-newline", ] as const; -type PromptFormatName = typeof PromptFormatNames[number]; +type PromptFormatName = (typeof PromptFormatNames)[number]; type PromptFormat = { - name: PromptFormatName; - makePrompt: (basePrompt: string) => string; - makeContinuation: (answer: string) => string; + name: PromptFormatName; + makePrompt: (basePrompt: string) => string; + makeContinuation: (answer: string) => string; }; type Option = { - label: string; - answer: string; + label: string; + answer: string; }; type ParsedSample = { - basePrompt: string; - answerLabel: string; - answer: string; - options: Option[]; + basePrompt: string; + answerLabel: string; + answer: string; + options: Option[]; }; type ScoreResult = { - score: number; - promptTokens: number; - continuationTokens: number; - usedInputTokens: number; + score: number; + promptTokens: number; + continuationTokens: number; + usedInputTokens: number; }; const promptFormats: PromptFormat[] = [ - { - name: "answer-colon-space", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, - makeContinuation: (answer) => answer, - }, - { - name: "answer-colon", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, - makeContinuation: (answer) => ` ${answer}`, - }, - { - name: "answer-newline", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, - makeContinuation: (answer) => answer, - }, + { + name: "answer-colon-space", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, + makeContinuation: (answer) => answer, + }, + { + name: "answer-colon", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, + makeContinuation: (answer) => ` ${answer}`, + }, + { + name: "answer-newline", + makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, + makeContinuation: (answer) => answer, + }, ]; function castPromptFormatName(raw: string): PromptFormatName { - for (const name of PromptFormatNames) { - if (raw === name) return name; - } - throw new Error(`Invalid promptFormat: ${raw}`); + for (const name of PromptFormatNames) { + if (raw === name) return name; + } + throw new Error(`Invalid promptFormat: ${raw}`); } function commonPrefixLength(left: number[], right: number[]): number { - const maxLength = Math.min(left.length, right.length); + const maxLength = Math.min(left.length, right.length); - for (let i = 0; i < maxLength; i++) { - if (left[i] !== right[i]) return i; - } + for (let i = 0; i < maxLength; i++) { + if (left[i] !== right[i]) return i; + } - return maxLength; + return maxLength; } -function predictTokenLogits(tfModel: tf.LayersModel, inputTensor: tf.Tensor2D): tf.Tensor3D { - const logits = tfModel.predict(inputTensor); - if (Array.isArray(logits)) { - throw new Error("Expected GPT model to return a single logits tensor"); - } - if (logits.rank !== 3) { - logits.dispose(); - throw new Error(`Expected GPT logits to have rank 3, got rank ${logits.rank}`); - } - return logits as tf.Tensor3D; +function predictTokenLogits( + tfModel: tf.LayersModel, + inputTensor: tf.Tensor2D, +): tf.Tensor3D { + const logits = tfModel.predict(inputTensor); + if (Array.isArray(logits)) { + throw new Error("Expected GPT model to return a single logits tensor"); + } + if (logits.rank !== 3) { + logits.dispose(); + throw new Error( + `Expected GPT logits to have rank 3, got rank ${logits.rank}`, + ); + } + return logits as tf.Tensor3D; } async function loadDataset(filePath: string, limit = -1): Promise { - const text = await fs.readFile(filePath, "utf-8"); - const samples = text - .split("<|endoftext|>") - .map((sample) => - sample - .replaceAll("<|startoftext|>", "") - .trim(), - ) - .filter((sample) => sample !== ""); - - return limit === -1 ? samples : samples.slice(0, limit); -} + const text = await fs.readFile(filePath, "utf-8"); + const samples = text + .split("<|endoftext|>") + .map((sample) => sample.replaceAll("<|startoftext|>", "").trim()) + .filter((sample) => sample !== ""); -function parseAnswerLine(line: string): { label: string; answer?: string } | undefined { - const match = line.trim().match(/^Answer:\s*([A-D])(?:\.\s*(.*))?$/i); - if (match === null) return undefined; - - const label = match[1].toUpperCase(); - const answerText = match[2]?.trim(); + return limit === -1 ? samples : samples.slice(0, limit); +} - return { - label, - answer: answerText === undefined || answerText === "" ? undefined : `${label}. ${answerText}`, - }; +function parseAnswerLine( + line: string, +): { label: string; answer?: string } | undefined { + const match = line.trim().match(/^Answer:\s*([A-D])(?:\.\s*(.*))?$/i); + if (match === null) return undefined; + + const label = match[1].toUpperCase(); + const answerText = match[2]?.trim(); + + return { + label, + answer: + answerText === undefined || answerText === "" + ? undefined + : `${label}. ${answerText}`, + }; } function parseOptionLine(line: string): Option | undefined { - const match = line.trim().match(/^([A-D])\.\s*(.+)$/i); - if (match === null) return undefined; - - const label = match[1].toUpperCase(); - return { - label, - answer: `${label}. ${match[2].trim()}`, - }; + const match = line.trim().match(/^([A-D])\.\s*(.+)$/i); + if (match === null) return undefined; + + const label = match[1].toUpperCase(); + return { + label, + answer: `${label}. ${match[2].trim()}`, + }; } function parseSample(sample: string): ParsedSample { - const lines = sample.split("\n"); - - let answerLabel = ""; - let answerFromLine: string | undefined; - const promptLines: string[] = []; - const options: Option[] = []; - - for (const line of lines) { - const answer = parseAnswerLine(line); - if (answer !== undefined) { - answerLabel = answer.label; - answerFromLine = answer.answer; - continue; - } - - const option = parseOptionLine(line); - if (option !== undefined) { - options.push(option); - } - - promptLines.push(line); + const lines = sample.split("\n"); + + let answerLabel = ""; + let answerFromLine: string | undefined; + const promptLines: string[] = []; + const options: Option[] = []; + + for (const line of lines) { + const answer = parseAnswerLine(line); + if (answer !== undefined) { + answerLabel = answer.label; + answerFromLine = answer.answer; + continue; } - const correctOption = options.find((option) => option.label === answerLabel); - if (correctOption === undefined) { - throw new Error(`Could not match answer label ${JSON.stringify(answerLabel)} to an option`); + const option = parseOptionLine(line); + if (option !== undefined) { + options.push(option); } - if (answerFromLine !== undefined && answerFromLine !== correctOption.answer) { - throw new Error( - `Answer line ${JSON.stringify(answerFromLine)} does not match option ${JSON.stringify(correctOption.answer)}`, - ); - } + promptLines.push(line); + } - const basePrompt = promptLines.join("\n").trim(); - return { - basePrompt, - answerLabel, - answer: correctOption.answer, - options, - }; + const correctOption = options.find((option) => option.label === answerLabel); + if (correctOption === undefined) { + throw new Error( + `Could not match answer label ${JSON.stringify(answerLabel)} to an option`, + ); + } + + if (answerFromLine !== undefined && answerFromLine !== correctOption.answer) { + throw new Error( + `Answer line ${JSON.stringify(answerFromLine)} does not match option ${JSON.stringify(correctOption.answer)}`, + ); + } + + const basePrompt = promptLines.join("\n").trim(); + return { + basePrompt, + answerLabel, + answer: correctOption.answer, + options, + }; } function validateOptions(options: Option[], expectedLabels: string[]): boolean { - if (options.length !== expectedLabels.length) return false; + if (options.length !== expectedLabels.length) return false; - const labels = options.map((option) => option.label); - return expectedLabels.every((label) => labels.includes(label)) - && new Set(labels).size === expectedLabels.length; + const labels = options.map((option) => option.label); + return ( + expectedLabels.every((label) => labels.includes(label)) && + new Set(labels).size === expectedLabels.length + ); } async function scoreContinuations( - tfModel: tf.LayersModel, - tokenizer: Tokenizer, - prompt: string, - continuations: string[], - contextLength: number, + tfModel: tf.LayersModel, + tokenizer: Tokenizer, + prompt: string, + continuations: string[], + contextLength: number, ): Promise { - const promptTokens = tokenizer.tokenize(prompt).toArray(); - const scoredInputs = continuations.map((continuation) => { - const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); - const continuationStart = commonPrefixLength(promptTokens, fullTokens); - const continuationTokens = fullTokens.length - continuationStart; - const inputTokens = fullTokens.slice(0, -1); - const offset = Math.max(0, inputTokens.length - contextLength); - const truncatedInputTokens = inputTokens.slice(offset); - - return { - fullTokens, - continuationStart, - continuationTokens, - offset, - truncatedInputTokens, - }; - }); - - const maxInputLength = scoredInputs.reduce( - (maxLength, scoredInput) => Math.max(maxLength, scoredInput.truncatedInputTokens.length), - 0, - ); - - if (maxInputLength === 0) { - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); - } - - const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ - ...truncatedInputTokens, - ...Array(maxInputLength - truncatedInputTokens.length).fill(0), - ]); - - const inputTensor = tf.tensor2d( - paddedInputs, - [paddedInputs.length, maxInputLength], - "int32", - ); - - const targetIndexes: number[][] = []; - const targetTokenIds: number[] = []; - const targetOwners: number[] = []; - - scoredInputs.forEach((scoredInput, batchIdx) => { - const { - fullTokens, - continuationStart, - offset, - truncatedInputTokens, - } = scoredInput; - - // Same ranking as HellaSwag's mean continuation cross-entropy: - // maximize mean log-probability instead of minimizing its negative. - for (let targetPos = continuationStart; targetPos < fullTokens.length; targetPos++) { - const targetToken = fullTokens[targetPos]; - const logitPos = targetPos - 1 - offset; - if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; - targetIndexes.push([batchIdx, logitPos]); - targetTokenIds.push(targetToken); - targetOwners.push(batchIdx); - } - }); - - if (targetIndexes.length === 0) { - inputTensor.dispose(); - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); - } - - const logits = predictTokenLogits(tfModel, inputTensor); - const targetLogProbs = tf.tidy(() => { - const targetIndexTensor = tf.tensor2d( - targetIndexes, - [targetIndexes.length, 2], - "int32", - ); - const targetTokenIndexTensor = tf.tensor2d( - targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), - [targetTokenIds.length, 2], - "int32", - ); - const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; - const logProbs = tf.logSoftmax(targetLogits, -1); - return tf.gatherND(logProbs, targetTokenIndexTensor); - }); - - const targetScores = await targetLogProbs.array() as number[]; - const scoreSums = Array(scoredInputs.length).fill(0) as number[]; - const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; - - targetScores.forEach((score, index) => { - const owner = targetOwners[index]; - scoreSums[owner] += score; - scoreCounts[owner]++; - }); + const promptTokens = tokenizer.tokenize(prompt).toArray(); + const scoredInputs = continuations.map((continuation) => { + const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); + const continuationStart = commonPrefixLength(promptTokens, fullTokens); + const continuationTokens = fullTokens.length - continuationStart; + const inputTokens = fullTokens.slice(0, -1); + const offset = Math.max(0, inputTokens.length - contextLength); + const truncatedInputTokens = inputTokens.slice(offset); - const results = scoredInputs.map((scoredInput, index) => ({ - score: scoreCounts[index] === 0 - ? Number.NEGATIVE_INFINITY - : scoreSums[index] / scoreCounts[index], - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, + return { + fullTokens, + continuationStart, + continuationTokens, + offset, + truncatedInputTokens, + }; + }); + + const maxInputLength = scoredInputs.reduce( + (maxLength, scoredInput) => + Math.max(maxLength, scoredInput.truncatedInputTokens.length), + 0, + ); + + if (maxInputLength === 0) { + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, })); + } + + const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ + ...truncatedInputTokens, + ...Array(maxInputLength - truncatedInputTokens.length).fill(0), + ]); + + const inputTensor = tf.tensor2d( + paddedInputs, + [paddedInputs.length, maxInputLength], + "int32", + ); + + const targetIndexes: number[][] = []; + const targetTokenIds: number[] = []; + const targetOwners: number[] = []; + + scoredInputs.forEach((scoredInput, batchIdx) => { + const { fullTokens, continuationStart, offset, truncatedInputTokens } = + scoredInput; + + // Same ranking as HellaSwag's mean continuation cross-entropy: + // maximize mean log-probability instead of minimizing its negative. + for ( + let targetPos = continuationStart; + targetPos < fullTokens.length; + targetPos++ + ) { + const targetToken = fullTokens[targetPos]; + const logitPos = targetPos - 1 - offset; + if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; + targetIndexes.push([batchIdx, logitPos]); + targetTokenIds.push(targetToken); + targetOwners.push(batchIdx); + } + }); + if (targetIndexes.length === 0) { inputTensor.dispose(); - logits.dispose(); - targetLogProbs.dispose(); - - return results; + return scoredInputs.map((scoredInput) => ({ + score: Number.NEGATIVE_INFINITY, + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + } + + const logits = predictTokenLogits(tfModel, inputTensor); + const targetLogProbs = tf.tidy(() => { + const targetIndexTensor = tf.tensor2d( + targetIndexes, + [targetIndexes.length, 2], + "int32", + ); + const targetTokenIndexTensor = tf.tensor2d( + targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), + [targetTokenIds.length, 2], + "int32", + ); + const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; + const logProbs = tf.logSoftmax(targetLogits, -1); + return tf.gatherND(logProbs, targetTokenIndexTensor); + }); + + const targetScores = (await targetLogProbs.array()) as number[]; + const scoreSums = Array(scoredInputs.length).fill(0) as number[]; + const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; + + targetScores.forEach((score, index) => { + const owner = targetOwners[index]; + scoreSums[owner] += score; + scoreCounts[owner]++; + }); + + const results = scoredInputs.map((scoredInput, index) => ({ + score: + scoreCounts[index] === 0 + ? Number.NEGATIVE_INFINITY + : scoreSums[index] / scoreCounts[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); + + inputTensor.dispose(); + logits.dispose(); + targetLogProbs.dispose(); + + return results; } async function benchmarkFullAnswers( - model: models.GPT, - tokenizer: Tokenizer, - dataset: string[], - format: PromptFormat, - contextLength: number, - savePath?: string, + model: models.GPT, + tokenizer: Tokenizer, + dataset: string[], + format: PromptFormat, + contextLength: number, + savePath?: string, ): Promise { - console.log(`=== FULL ANSWER LOGPROB BENCHMARK (${format.name}) ===`); - console.log(`Context length: ${contextLength}`); - - const tfModel = model.extract(); - - let correct = 0; - let total = 0; - const labels = ["A", "B", "C", "D"]; - const confusion: Record> = Object.fromEntries( - labels.map((label) => [ - label, - Object.fromEntries(labels.map((otherLabel) => [otherLabel, 0])), - ]), - ); - - type PredictionLog = { - predicted: string; - predictedAnswer: string; - answer: string; - answerText: string; - correct: boolean; - scores: Record; - promptTokens: number; - continuationTokens: Record; - usedInputTokens: Record; - }; - - const logs: PredictionLog[] = []; - const start = Date.now(); - - for (const sample of dataset) { - let parsed: ParsedSample; - try { - parsed = parseSample(sample); - } catch (error) { - console.log("Invalid sample:", error instanceof Error ? error.message : error); - continue; - } - - if (!validateOptions(parsed.options, labels)) { - console.log("Invalid options:", parsed.options.map((option) => option.label).join(", ")); - continue; - } - - const prompt = format.makePrompt(parsed.basePrompt); - const results = await scoreContinuations( - tfModel, - tokenizer, - prompt, - parsed.options.map((option) => format.makeContinuation(option.answer)), - contextLength, - ); - - const scores = results.map((result) => result.score); - let bestIdx = 0; - for (let i = 1; i < scores.length; i++) { - if (scores[i] > scores[bestIdx]) bestIdx = i; - } - - const predicted = parsed.options[bestIdx]; - if (predicted.label === parsed.answerLabel) correct++; - total++; - - if (confusion[parsed.answerLabel]?.[predicted.label] === undefined) { - throw new Error( - `Unexpected confusion matrix key: answer=${parsed.answerLabel}, predicted=${predicted.label}`, - ); - } - confusion[parsed.answerLabel][predicted.label]++; - - logs.push({ - predicted: predicted.label, - predictedAnswer: predicted.answer, - answer: parsed.answerLabel, - answerText: parsed.answer, - correct: predicted.label === parsed.answerLabel, - scores: Object.fromEntries( - parsed.options.map((option, i) => [option.label, scores[i]]), - ), - promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, - continuationTokens: Object.fromEntries( - parsed.options.map((option, i) => [option.label, results[i].continuationTokens]), - ), - usedInputTokens: Object.fromEntries( - parsed.options.map((option, i) => [option.label, results[i].usedInputTokens]), - ), - }); - - if (total % 50 === 0) { - console.log(`Processed ${total} samples...`); - } + console.log(`=== FULL ANSWER LOGPROB BENCHMARK (${format.name}) ===`); + console.log(`Context length: ${contextLength}`); + + const tfModel = model.extract(); + + let correct = 0; + let total = 0; + const labels = ["A", "B", "C", "D"]; + const confusion: Record> = Object.fromEntries( + labels.map((label) => [ + label, + Object.fromEntries(labels.map((otherLabel) => [otherLabel, 0])), + ]), + ); + + type PredictionLog = { + predicted: string; + predictedAnswer: string; + answer: string; + answerText: string; + correct: boolean; + scores: Record; + promptTokens: number; + continuationTokens: Record; + usedInputTokens: Record; + }; + + const logs: PredictionLog[] = []; + const start = Date.now(); + + for (const sample of dataset) { + let parsed: ParsedSample; + try { + parsed = parseSample(sample); + } catch (error) { + console.log( + "Invalid sample:", + error instanceof Error ? error.message : error, + ); + continue; } - if (total === 0) { - throw new Error("No valid samples were evaluated"); + if (!validateOptions(parsed.options, labels)) { + console.log( + "Invalid options:", + parsed.options.map((option) => option.label).join(", "), + ); + continue; } - const accuracy = correct / total; - const duration = ((Date.now() - start) / 1000).toFixed(2); - - console.log("\n========================="); - console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); - console.log(`Time: ${duration}s`); - console.log("=========================\n"); + const prompt = format.makePrompt(parsed.basePrompt); + const results = await scoreContinuations( + tfModel, + tokenizer, + prompt, + parsed.options.map((option) => format.makeContinuation(option.answer)), + contextLength, + ); - console.log("Confusion Matrix:"); - console.table(confusion); + const scores = results.map((result) => result.score); + let bestIdx = 0; + for (let i = 1; i < scores.length; i++) { + if (scores[i] > scores[bestIdx]) bestIdx = i; + } - console.log("\nPer-class accuracy:"); - for (const cls of labels) { - const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); - const correctCls = confusion[cls][cls]; - const acc = totalCls ? (correctCls / totalCls) * 100 : 0; + const predicted = parsed.options[bestIdx]; + if (predicted.label === parsed.answerLabel) correct++; + total++; - console.log(`${cls}: ${acc.toFixed(2)}%`); + if (confusion[parsed.answerLabel]?.[predicted.label] === undefined) { + throw new Error( + `Unexpected confusion matrix key: answer=${parsed.answerLabel}, predicted=${predicted.label}`, + ); } + confusion[parsed.answerLabel][predicted.label]++; + + logs.push({ + predicted: predicted.label, + predictedAnswer: predicted.answer, + answer: parsed.answerLabel, + answerText: parsed.answer, + correct: predicted.label === parsed.answerLabel, + scores: Object.fromEntries( + parsed.options.map((option, i) => [option.label, scores[i]]), + ), + promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, + continuationTokens: Object.fromEntries( + parsed.options.map((option, i) => [ + option.label, + results[i].continuationTokens, + ]), + ), + usedInputTokens: Object.fromEntries( + parsed.options.map((option, i) => [ + option.label, + results[i].usedInputTokens, + ]), + ), + }); - if (savePath) { - await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); - console.log(`Saved results to ${savePath}`); + if (total % 50 === 0) { + console.log(`Processed ${total} samples...`); } + } - return accuracy; -} - -async function main() { - const args = parse({ - modelPath: { type: String }, - testPath: { type: String }, - maxSamples: { type: Number, optional: true, defaultValue: 100 }, - savePath: { type: String, optional: true }, - compareFormats: { type: Boolean, optional: true, defaultValue: false }, - promptFormat: { - type: (raw: string) => castPromptFormatName(raw), - optional: true, - defaultValue: "answer-colon-space", - }, - contextLength: { type: Number, optional: true }, - help: { type: Boolean, optional: true }, - }); + if (total === 0) { + throw new Error("No valid samples were evaluated"); + } - console.log("Loading tokenizer..."); - const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + const accuracy = correct / total; + const duration = ((Date.now() - start) / 1000).toFixed(2); - console.log("Loading model..."); - const model = await loadModelFromDisk(args.modelPath); + console.log("\n========================="); + console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); + console.log(`Time: ${duration}s`); + console.log("=========================\n"); - if (!(model instanceof models.GPT)) { - throw new Error("Model must be GPT"); - } + console.log("Confusion Matrix:"); + console.table(confusion); - console.log("Loading dataset..."); - const dataset = await loadDataset(args.testPath, args.maxSamples); + console.log("\nPer-class accuracy:"); + for (const cls of labels) { + const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); + const correctCls = confusion[cls][cls]; + const acc = totalCls ? (correctCls / totalCls) * 100 : 0; - console.log(`Loaded ${dataset.length} samples`); + console.log(`${cls}: ${acc.toFixed(2)}%`); + } - const contextLength = args.contextLength ?? model.config.contextLength; - const formats = args.compareFormats - ? promptFormats - : promptFormats.filter((format) => format.name === args.promptFormat); + if (savePath) { + await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); + console.log(`Saved results to ${savePath}`); + } - for (const format of formats) { - const savePath = - args.savePath === undefined || formats.length === 1 - ? args.savePath - : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + return accuracy; +} - await benchmarkFullAnswers(model, tokenizer, dataset, format, contextLength, savePath); - } +async function main() { + const args = parse({ + modelPath: { type: String }, + testPath: { type: String }, + maxSamples: { type: Number, optional: true, defaultValue: 100 }, + savePath: { type: String, optional: true }, + compareFormats: { type: Boolean, optional: true, defaultValue: false }, + promptFormat: { + type: (raw: string) => castPromptFormatName(raw), + optional: true, + defaultValue: "answer-colon-space", + }, + contextLength: { type: Number, optional: true }, + help: { type: Boolean, optional: true }, + }); + + console.log("Loading tokenizer..."); + const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); + + console.log("Loading model..."); + const model = await loadModelFromDisk(args.modelPath); + + if (!(model instanceof models.GPT)) { + throw new Error("Model must be GPT"); + } + + console.log("Loading dataset..."); + const dataset = await loadDataset(args.testPath, args.maxSamples); + + console.log(`Loaded ${dataset.length} samples`); + + const contextLength = args.contextLength ?? model.config.contextLength; + const formats = args.compareFormats + ? promptFormats + : promptFormats.filter((format) => format.name === args.promptFormat); + + for (const format of formats) { + const savePath = + args.savePath === undefined || formats.length === 1 + ? args.savePath + : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + + await benchmarkFullAnswers( + model, + tokenizer, + dataset, + format, + contextLength, + savePath, + ); + } - console.log("Done."); + console.log("Done."); } main().catch(console.error); diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts index fbbdb7b38..165beb589 100644 --- a/cli/src/measure_memorization_gpt2.ts +++ b/cli/src/measure_memorization_gpt2.ts @@ -23,7 +23,7 @@ interface Args { } const DecodingStrategies = ["top-k", "greedy"] as const; -type DecodingStrategy = typeof DecodingStrategies[number]; +type DecodingStrategy = (typeof DecodingStrategies)[number]; type PromptResult = { recordIndex: number; @@ -53,7 +53,9 @@ function parseIntegerList(raw: string): number[] { .filter((v) => !Number.isNaN(v)); if (values.length === 0 || values.some((v) => v <= 0)) { - throw new Error("promptLengths must be a comma-separated list of positive integers"); + throw new Error( + "promptLengths must be a comma-separated list of positive integers", + ); } return values; @@ -75,12 +77,18 @@ function seededRandom(seed: number): () => number { }; } -function randomInt(random: () => number, minInclusive: number, maxInclusive: number): number { +function randomInt( + random: () => number, + minInclusive: number, + maxInclusive: number, +): number { if (maxInclusive < minInclusive) { throw new Error("invalid random integer range"); } - return minInclusive + Math.floor(random() * (maxInclusive - minInclusive + 1)); + return ( + minInclusive + Math.floor(random() * (maxInclusive - minInclusive + 1)) + ); } async function loadRecords(filePath: string, limit: number): Promise { @@ -157,7 +165,8 @@ function bleu1to4(reference: number[], candidate: number[]): number { ? 1 : Math.exp(1 - reference.length / candidate.length); const geometricMean = Math.exp( - precisions.reduce((sum, precision) => sum + Math.log(precision), 0) / precisions.length, + precisions.reduce((sum, precision) => sum + Math.log(precision), 0) / + precisions.length, ); return brevityPenalty * geometricMean; @@ -191,21 +200,15 @@ async function sampleGenerateGPT2( const nextTokenTensor = tf.tidy(() => { const last = logits.slice([0, modelInput.length - 1, 0], [1, 1, -1]); const scaled = last.squeeze().div(temperature); - const { values: topKLogits, indices: topKTokens } = tf.topk( - scaled, - topK, - ); + const { values: topKLogits, indices: topKTokens } = tf.topk(scaled, topK); if (decodingStrategy === "greedy") { return topKTokens.gather(tf.scalar(0, "int32")).squeeze(); } - const sampledIndex = tf.multinomial( - topKLogits.expandDims(0), - 1, - seed + i, - false, - ).squeeze(); + const sampledIndex = tf + .multinomial(topKLogits.expandDims(0), 1, seed + i, false) + .squeeze(); return topKTokens.gather(sampledIndex).squeeze(); }); @@ -226,16 +229,17 @@ async function sampleGenerateGPT2( function summarize(results: PromptResult[]) { const byPromptLength = new Map(); for (const result of results) { - byPromptLength.set( - result.promptLength, - [...(byPromptLength.get(result.promptLength) ?? []), result], - ); + byPromptLength.set(result.promptLength, [ + ...(byPromptLength.get(result.promptLength) ?? []), + result, + ]); } const summarizeGroup = (group: PromptResult[]) => ({ count: group.length, exactMatchRate: group.filter((r) => r.exactMatch).length / group.length, - bleuMemorizationRate: group.filter((r) => r.memorizedByBleu).length / group.length, + bleuMemorizationRate: + group.filter((r) => r.memorizedByBleu).length / group.length, averageBleu: group.reduce((sum, r) => sum + r.bleu, 0) / group.length, }); @@ -253,30 +257,80 @@ function summarize(results: PromptResult[]) { async function main() { const args = parse( { - modelPath: { type: String, description: "Path to a saved Disco GPT model.json" }, - dataPath: { type: String, description: "Path to records/canaries text file" }, - maxRecords: { type: Number, description: "Maximum records to evaluate; -1 for all", defaultValue: 100 }, - promptLengths: { type: String, description: "Comma-separated prompt lengths", defaultValue: "10,50,100,200,500" }, - suffixLength: { type: Number, description: "Number of suffix tokens to generate and compare", defaultValue: 50 }, - bleuThreshold: { type: Number, description: "BLEU threshold for approximate memorization", defaultValue: 0.75 }, + modelPath: { + type: String, + description: "Path to a saved Disco GPT model.json", + }, + dataPath: { + type: String, + description: "Path to records/canaries text file", + }, + maxRecords: { + type: Number, + description: "Maximum records to evaluate; -1 for all", + defaultValue: 100, + }, + promptLengths: { + type: String, + description: "Comma-separated prompt lengths", + defaultValue: "10,50,100,200,500", + }, + suffixLength: { + type: Number, + description: "Number of suffix tokens to generate and compare", + defaultValue: 50, + }, + bleuThreshold: { + type: Number, + description: "BLEU threshold for approximate memorization", + defaultValue: 0.75, + }, decodingStrategy: { type: (raw: string) => castDecodingStrategy(raw), description: "Generation strategy: top-k or greedy", defaultValue: "top-k", }, - temperature: { type: Number, description: "Generation temperature used with top-k sampling", defaultValue: 0.8 }, - topK: { type: Number, description: "Number of most likely tokens considered for top-k sampling", defaultValue: 50 }, - seed: { type: Number, description: "Random seed for choosing record split positions", defaultValue: 42 }, - logEvery: { type: Number, description: "Print progress every N records; set 0 to disable per-record progress logs", defaultValue: 1 }, - savePath: { type: String, description: "Optional JSON output path", optional: true }, - help: { type: Boolean, optional: true, alias: "h", description: "Prints this usage guide" }, + temperature: { + type: Number, + description: "Generation temperature used with top-k sampling", + defaultValue: 0.8, + }, + topK: { + type: Number, + description: + "Number of most likely tokens considered for top-k sampling", + defaultValue: 50, + }, + seed: { + type: Number, + description: "Random seed for choosing record split positions", + defaultValue: 42, + }, + logEvery: { + type: Number, + description: + "Print progress every N records; set 0 to disable per-record progress logs", + defaultValue: 1, + }, + savePath: { + type: String, + description: "Optional JSON output path", + optional: true, + }, + help: { + type: Boolean, + optional: true, + alias: "h", + description: "Prints this usage guide", + }, }, { helpArg: "help", headerContentSections: [ { header: "GPT-2 Unintended Memorization", - content: "Measures extractable memorization via greedy suffix generation.", + content: + "Measures extractable memorization via greedy suffix generation.", }, ], }, @@ -305,7 +359,9 @@ async function main() { console.log(`Loaded ${records.length} records`); console.log("Tokenizing records..."); - const tokenizedRecords = records.map((record) => tokenizer.tokenize(record).toArray()); + const tokenizedRecords = records.map((record) => + tokenizer.tokenize(record).toArray(), + ); const tokenLengths = tokenizedRecords.map((ids) => ids.length); const requiredTokensByPromptLength = Object.fromEntries( promptLengths.map((promptLength) => [ @@ -322,7 +378,10 @@ async function main() { ]), ); console.log("Token length stats:", summarizeTokenLengths(tokenLengths)); - console.log("Eligible records by prompt length:", eligibleRecordsByPromptLength); + console.log( + "Eligible records by prompt length:", + eligibleRecordsByPromptLength, + ); console.log("Starting memorization evaluation..."); const results: PromptResult[] = []; @@ -331,7 +390,11 @@ async function main() { promptLengths.map((promptLength) => [promptLength, 0]), ); - for (let recordIndex = 0; recordIndex < tokenizedRecords.length; recordIndex++) { + for ( + let recordIndex = 0; + recordIndex < tokenizedRecords.length; + recordIndex++ + ) { const ids = tokenizedRecords[recordIndex]; const eligiblePromptLengths = promptLengths.filter( (promptLength) => ids.length >= promptLength + args.suffixLength + 1, @@ -344,7 +407,10 @@ async function main() { if (shouldLogRecord) { console.log( - `Record ${recordIndex + 1}/${tokenizedRecords.length}: ${ids.length} tokens, eligible prompt lengths: ${eligiblePromptLengths.length > 0 ? eligiblePromptLengths.join(",") : "none" + `Record ${recordIndex + 1}/${tokenizedRecords.length}: ${ids.length} tokens, eligible prompt lengths: ${ + eligiblePromptLengths.length > 0 + ? eligiblePromptLengths.join(",") + : "none" }`, ); } @@ -358,7 +424,8 @@ async function main() { if (eligiblePromptLengths.length === 0) { if (shouldLogRecord) { console.log( - `Skipping record ${recordIndex + 1}; needs at least ${Math.min(...promptLengths) + args.suffixLength + 1 + `Skipping record ${recordIndex + 1}; needs at least ${ + Math.min(...promptLengths) + args.suffixLength + 1 } tokens for the shortest prompt/suffix setting.`, ); } @@ -392,7 +459,10 @@ async function main() { args.topK, args.seed + recordIndex + promptLength, ); - const generatedSuffix = generated.slice(prompt.length, prompt.length + args.suffixLength); + const generatedSuffix = generated.slice( + prompt.length, + prompt.length + args.suffixLength, + ); console.log("================================"); console.log("PROMPT LENGTH:", promptLength); diff --git a/discojs/src/aggregator/mean.ts b/discojs/src/aggregator/mean.ts index ba45ef848..080ff6a03 100644 --- a/discojs/src/aggregator/mean.ts +++ b/discojs/src/aggregator/mean.ts @@ -18,7 +18,9 @@ export class MeanAggregator extends MultiRoundAggregator { } override _add(nodeId: client.NodeID, contribution: WeightsContainer): void { - const previous = this.contributions.getIn([0, nodeId]) as WeightsContainer | undefined; + const previous = this.contributions.getIn([0, nodeId]) as + | WeightsContainer + | undefined; this.log( this.contributions.hasIn([0, nodeId]) ? AggregationStep.UPDATE @@ -41,7 +43,8 @@ export class MeanAggregator extends MultiRoundAggregator { const contributions = Array.from(currentContributions.values()); let summed = contributions[0]?.map((weight) => weight.clone()); - if (summed === undefined) throw new Error("aggregating without any contribution"); + if (summed === undefined) + throw new Error("aggregating without any contribution"); try { for (const contribution of contributions.slice(1)) { diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 3fc5451cd..9577397a2 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -165,12 +165,13 @@ export abstract class Client extends EventEmitter<{ `[${shortenId(this.ownId)}] received EnoughParticipants message from server`, ); // Emit the last status emitted before waiting if defined - if (this.#previousStatus !== undefined) this.emit("status", this.#previousStatus) - this.nbOfParticipants = event.nbOfParticipants - this.promiseForMoreParticipants = undefined - resolve() - }) - }) + if (this.#previousStatus !== undefined) + this.emit("status", this.#previousStatus); + this.nbOfParticipants = event.nbOfParticipants; + this.promiseForMoreParticipants = undefined; + resolve(); + }); + }); } protected async waitForParticipantsIfNeeded(): Promise { @@ -197,11 +198,15 @@ export abstract class Client extends EventEmitter<{ } url.pathname += `tasks/${this.task.id}/model.json`; - debug("fetching latest model from server at %0 for task %1...", url.href, this.task.id) + debug( + "fetching latest model from server at %0 for task %1...", + url.href, + this.task.id, + ); const response = await fetch(url); - if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`) - else debug("response ok, decoding model...") + if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`); + else debug("response ok, decoding model..."); const encoded = new Uint8Array(await response.arrayBuffer()); return await serialization.model.decode(encoded); diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index fda02f5f2..90ca25c3d 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -121,12 +121,13 @@ export class WebSocketServer static async connect( url: URL, validateReceived: (msg: unknown) => msg is Message, - validateSent: (msg: Message) => boolean): Promise { + validateSent: (msg: Message) => boolean, + ): Promise { const ws = new WebSocket(url, { // Federated GPT updates can exceed the default ws payload limit. maxPayload: 1024 * 1024 * 1024, - }) - ws.binaryType = 'arraybuffer' + }); + ws.binaryType = "arraybuffer"; const server: WebSocketServer = new WebSocketServer(ws, validateSent); @@ -146,16 +147,23 @@ export class WebSocketServer }; ws.onclose = (event) => { - debug("websocket closed: code=%o reason=%o wasClean=%o", event.code, event.reason, event.wasClean) - } + debug( + "websocket closed: code=%o reason=%o wasClean=%o", + event.code, + event.reason, + event.wasClean, + ); + }; return await new Promise((resolve, reject) => { ws.onerror = (err: WebSocket.ErrorEvent) => { - debug("websocket error while connecting/receiving: %o", err.message) - reject(new Error(`Server unreachable: ${err.message}`)) - } - ws.onopen = () => { resolve(server) } - }) + debug("websocket error while connecting/receiving: %o", err.message); + reject(new Error(`Server unreachable: ${err.message}`)); + }; + ws.onopen = () => { + resolve(server); + }; + }); } disconnect(): Promise { diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index ccd13d66b..1f5fbddb7 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -93,12 +93,15 @@ export class FederatedClient extends Client<"federated"> { this.nbOfParticipants = nbOfParticipants; // Upon connecting, the server answers with a boolean // which indicates whether there are enough participants or not - debug(`[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants) + debug( + `[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, + this.waitingForMoreParticipants, + ); if (payload != null) { - const latestWeights = serialization.weights.decode(payload) - model.weights = latestWeights + const latestWeights = serialization.weights.decode(payload); + model.weights = latestWeights; } - return model + return model; } /** @@ -145,17 +148,21 @@ export class FederatedClient extends Client<"federated"> { this.saveAndEmit("updating model"); // Send our local contribution to the server // and receive the server global update for this round as an answer to our contribution - const payloadToServer = this.aggregator - .makePayloads(weights) - .get(SERVER_NODE_ID); - if (payloadToServer === undefined) - throw new Error("aggregator didn't make a payload for the server"); + const payloadToServer = this.aggregator + .makePayloads(weights) + .get(SERVER_NODE_ID); + if (payloadToServer === undefined) + throw new Error("aggregator didn't make a payload for the server"); const round = this.aggregator.round; { - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before encode`); + debugProcessMemory( + `[${shortenId(this.ownId)}] round ${round} before encode`, + ); const payload = await serialization.weights.encode(payloadToServer); - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after encode`); + debugProcessMemory( + `[${shortenId(this.ownId)}] round ${round} after encode`, + ); debug( "[%s] encoded payload for round %d byteLength=%d", shortenId(this.ownId), @@ -171,21 +178,33 @@ export class FederatedClient extends Client<"federated"> { // Need to await the resulting global model right after sending our local contribution // to make sure we don't miss it - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} before send`); + debugProcessMemory( + `[${shortenId(this.ownId)}] round ${round} before send`, + ); this.server.send(msg); - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after send`); + debugProcessMemory( + `[${shortenId(this.ownId)}] round ${round} after send`, + ); } - debug(`[${shortenId(this.ownId)}] sent its local update to the server for round ${round}`); - debug(`[${shortenId(this.ownId)}] is waiting for server update for round ${round + 1}`); + debug( + `[${shortenId(this.ownId)}] sent its local update to the server for round ${round}`, + ); + debug( + `[${shortenId(this.ownId)}] is waiting for server update for round ${round + 1}`, + ); const { payload: payloadFromServer, round: serverRound, - nbOfParticipants - } = await waitMessage( this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update - this.nbOfParticipants = nbOfParticipants // Save the current participants - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after receive`); + nbOfParticipants, + } = await waitMessage(this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update + this.nbOfParticipants = nbOfParticipants; // Save the current participants + debugProcessMemory( + `[${shortenId(this.ownId)}] round ${round} after receive`, + ); const serverResult = serialization.weights.decode(payloadFromServer); - debugProcessMemory(`[${shortenId(this.ownId)}] round ${round} after decode server payload`); + debugProcessMemory( + `[${shortenId(this.ownId)}] round ${round} after decode server payload`, + ); this.aggregator.setRound(serverRound); return serverResult; diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts index 32ec504de..8b848f544 100644 --- a/discojs/src/client/federated/messages.ts +++ b/discojs/src/client/federated/messages.ts @@ -19,12 +19,12 @@ export type MessageFederated = | EnoughParticipants; export interface NewFederatedNodeInfo { - type: type.NewFederatedNodeInfo - id: NodeID - waitForMoreParticipants: boolean + type: type.NewFederatedNodeInfo; + id: NodeID; + waitForMoreParticipants: boolean; payload?: serialization.Encoded | null; - round: number - nbOfParticipants: number + round: number; + nbOfParticipants: number; } export interface SendPayload { diff --git a/discojs/src/client/local_client.ts b/discojs/src/client/local_client.ts index 2d87292d4..e6a89e68e 100644 --- a/discojs/src/client/local_client.ts +++ b/discojs/src/client/local_client.ts @@ -11,7 +11,9 @@ export class LocalClient extends Client<"local"> { } // Return clones so the trainer can dispose the communication result without // disposing tensors owned by the model. - override onRoundEndCommunication(weights: WeightsContainer): Promise { + override onRoundEndCommunication( + weights: WeightsContainer, + ): Promise { return Promise.resolve(weights.map((weight) => weight.clone())); } } diff --git a/discojs/src/default_tasks/centralized_gpt2_finetune.ts b/discojs/src/default_tasks/centralized_gpt2_finetune.ts index 833a1622a..ea6a8df5a 100644 --- a/discojs/src/default_tasks/centralized_gpt2_finetune.ts +++ b/discojs/src/default_tasks/centralized_gpt2_finetune.ts @@ -4,13 +4,15 @@ import { Tokenizer, models, serialization } from "../index.js"; export const centralizedGPT2FineTune: TaskProvider<"text", "local"> = { async getTask() { return { - id: 'centralized-gpt2-finetune', + id: "centralized-gpt2-finetune", dataType: "text", displayInformation: { title: "GPT Privacy-Preserving Fine-tuning", summary: { - preview: 'Fine-tune a pre-trained GPT model collaboratively and privately.', - overview: "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning." + preview: + "Fine-tune a pre-trained GPT model collaboratively and privately.", + overview: + "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning.", }, model: [ "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", @@ -18,52 +20,58 @@ export const centralizedGPT2FineTune: TaskProvider<"text", "local"> = { "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", "Context length is kept at 1024 to match the pre-trained model, with batch size at 1.", ].join(" "), - dataFormatInformation: 'You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.', + dataFormatInformation: + "You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.", dataExample: "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", }, trainingInformation: { - scheme: 'local', - aggregationStrategy: 'mean', + scheme: "local", + aggregationStrategy: "mean", epochs: 1, - validationSplit: 0.1, + validationSplit: 0.1, roundDuration: 1, // Last context segment may be shorter than context length, so it will be dropped (TODO: implement padding to avoid this) batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), contextLength: 512, - tensorBackend: 'gpt' - } - } + tensorBackend: "gpt", + }, + }; }, async getModel() { // Load the pre-trained ONNX-converted model from Google Cloud Storage // The model should be in DiscoJS serialization format (created by onnx-converter) - const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; + const modelUrl = + "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; try { const response = await fetch(modelUrl); if (!response.ok) { throw new Error(`HTTP ${response.status}: ${response.statusText}`); } - + const arrayBuffer = await response.arrayBuffer(); const encodedData = new Uint8Array(arrayBuffer); - + const model = await serialization.model.decode(encodedData); - + if (!(model instanceof models.GPT)) { throw new Error("Loaded model is not a GPT model"); } - console.log("Successfully loaded pre-trained GPT model from Google Cloud Storage"); - + console.log( + "Successfully loaded pre-trained GPT model from Google Cloud Storage", + ); + return model; } catch (error) { console.error("Failed to load model from Google Cloud Storage:", error); - throw new Error(`Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`); + throw new Error( + `Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`, + ); } }, -} +}; diff --git a/discojs/src/default_tasks/index.ts b/discojs/src/default_tasks/index.ts index 4d359574d..1ca9eba9c 100644 --- a/discojs/src/default_tasks/index.ts +++ b/discojs/src/default_tasks/index.ts @@ -1,9 +1,9 @@ -export { cifar10 } from './cifar10.js' -export { lusCovid } from './lus_covid.js' -export { mnist } from './mnist.js' -export { simpleFace } from './simple_face.js' -export { titanic } from './titanic.js' -export { wikitext } from './wikitext.js' -export { tinderDog } from './tinder_dog.js' -export { privacyrun } from './privacyrun.js' -export { centralizedGPT2FineTune } from './centralized_gpt2_finetune.js' +export { cifar10 } from "./cifar10.js"; +export { lusCovid } from "./lus_covid.js"; +export { mnist } from "./mnist.js"; +export { simpleFace } from "./simple_face.js"; +export { titanic } from "./titanic.js"; +export { wikitext } from "./wikitext.js"; +export { tinderDog } from "./tinder_dog.js"; +export { privacyrun } from "./privacyrun.js"; +export { centralizedGPT2FineTune } from "./centralized_gpt2_finetune.js"; diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index d1bb23863..3bd0cc984 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -4,13 +4,15 @@ import { Tokenizer, models, serialization } from "../index.js"; export const privacyrun: TaskProvider<"text", "federated"> = { async getTask() { return { - id: 'privacyrun', + id: "privacyrun", dataType: "text", displayInformation: { title: "GPT Privacy-Preserving Fine-tuning", summary: { - preview: 'Fine-tune a pre-trained GPT model collaboratively and privately.', - overview: "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning." + preview: + "Fine-tune a pre-trained GPT model collaboratively and privately.", + overview: + "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning.", }, model: [ "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", @@ -18,55 +20,61 @@ export const privacyrun: TaskProvider<"text", "federated"> = { "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", "Context length is kept at 1024 to match the pre-trained model, with batch size at 1.", ].join(" "), - dataFormatInformation: 'You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.', + dataFormatInformation: + "You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.", dataExample: "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", }, trainingInformation: { - scheme: 'federated', - aggregationStrategy: 'mean', + scheme: "federated", + aggregationStrategy: "mean", minNbOfParticipants: 2, epochs: 1, - validationSplit: 0.1, + validationSplit: 0.1, roundDuration: 1, // Last context segment may be shorter than context length, so it will be dropped (TODO: implement padding to avoid this) batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), // contextLength: 1024, contextLength: 512, - tensorBackend: 'gpt' - } - } + tensorBackend: "gpt", + }, + }; }, async getModel() { // Load the pre-trained ONNX-converted model from Google Cloud Storage // The model should be in DiscoJS serialization format (created by onnx-converter) // const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model.json"; - - const modelUrl = "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; + + const modelUrl = + "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; try { const response = await fetch(modelUrl); if (!response.ok) { throw new Error(`HTTP ${response.status}: ${response.statusText}`); } - + const arrayBuffer = await response.arrayBuffer(); const encodedData = new Uint8Array(arrayBuffer); - + const model = await serialization.model.decode(encodedData); - + if (!(model instanceof models.GPT)) { throw new Error("Loaded model is not a GPT model"); } - console.log("Successfully loaded pre-trained GPT model from Google Cloud Storage"); - + console.log( + "Successfully loaded pre-trained GPT model from Google Cloud Storage", + ); + return model; } catch (error) { console.error("Failed to load model from Google Cloud Storage:", error); - throw new Error(`Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`); + throw new Error( + `Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`, + ); } }, -} +}; diff --git a/discojs/src/models/gpt/config.ts b/discojs/src/models/gpt/config.ts index c01b7aacb..e1fdd3b68 100644 --- a/discojs/src/models/gpt/config.ts +++ b/discojs/src/models/gpt/config.ts @@ -8,32 +8,32 @@ type GPTModelType = | "gpt-nano"; export type GPTConfig = { - lr: number - contextLength: number - vocabSize?: number - modelType: GPTModelType - evaluate?: boolean - maxEvalBatches?: number - evaluateEvery?: number - maxIter?: number - weightDecay?: number - verbose?: 0 | 1 - debug?: boolean - attnDrop?: number - residDrop?: number - embdDrop?: number - nLayer?: number - nHead?: number - nEmbd?: number - seed?: number, -} + lr: number; + contextLength: number; + vocabSize?: number; + modelType: GPTModelType; + evaluate?: boolean; + maxEvalBatches?: number; + evaluateEvery?: number; + maxIter?: number; + weightDecay?: number; + verbose?: 0 | 1; + debug?: boolean; + attnDrop?: number; + residDrop?: number; + embdDrop?: number; + nLayer?: number; + nHead?: number; + nEmbd?: number; + seed?: number; +}; export type GoldfishLossConfig = { - enabled: boolean - k: number - h: number - padTokenId?: number -} + enabled: boolean; + k: number; + h: number; + padTokenId?: number; +}; // for a benchmark of performance, see https://github.com/epfml/disco/pull/659 export const DefaultGPTConfig: Required = { lr: 0.001, diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index 8f32c2efb..e75dda4e1 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -14,8 +14,12 @@ import { BatchLogs, Model, EpochLogs } from "../index.js"; import { GPTModel } from "./model.js"; import evaluate from "./evaluate.js"; -import { DefaultGPTConfig, DefaultGenerationConfig } from './config.js' -import type { GoldfishLossConfig, GPTConfig, GenerationConfig } from './config.js' +import { DefaultGPTConfig, DefaultGenerationConfig } from "./config.js"; +import type { + GoldfishLossConfig, + GPTConfig, + GenerationConfig, +} from "./config.js"; const debug = createDebug("discojs:models:gpt"); @@ -109,7 +113,10 @@ export class GPT extends Model<"text"> { break; } - const batchLogs = await this.#runBatch(next.value, ++this.#iterationCount); + const batchLogs = await this.#runBatch( + next.value, + ++this.#iterationCount, + ); yield batchLogs; batchesLogs = batchesLogs.push(batchLogs); @@ -286,18 +293,17 @@ export class GPT extends Model<"text"> { } static deserialize(data: GPTSerialization): Model<"text"> { - - debug("GPT model deserialization started") + debug("GPT model deserialization started"); const config = data.config; const model = new GPT(config); - debug("GPT model config initialized: %O", config) + debug("GPT model config initialized: %O", config); model.weights = data.weights; - debug("GPT model weights initialized") + debug("GPT model weights initialized"); return model; } diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 70d6ef8aa..cc9506537 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -1,11 +1,11 @@ import createDebug from "debug"; import * as tf from "@tensorflow/tfjs"; -import type { GoldfishLossConfig, GPTConfig } from './config.js' -import { getModelSizes, DefaultGPTConfig } from './config.js' -import { getCustomAdam, clipByGlobalNormObj } from './optimizers.js' -import evaluate from './evaluate.js' -import { GPTArchitecture } from './layers.js' +import type { GoldfishLossConfig, GPTConfig } from "./config.js"; +import { getModelSizes, DefaultGPTConfig } from "./config.js"; +import { getCustomAdam, clipByGlobalNormObj } from "./optimizers.js"; +import evaluate from "./evaluate.js"; +import { GPTArchitecture } from "./layers.js"; const debug = createDebug("discojs:models:gpt:model"); @@ -38,9 +38,9 @@ export declare abstract class Dataset { * */ export class GPTModel extends tf.LayersModel { - protected readonly config: Required - #debugLabel?: string - #goldfishLoss?: GoldfishLossConfig + protected readonly config: Required; + #debugLabel?: string; + #goldfishLoss?: GoldfishLossConfig; constructor( partialConfig?: Partial, @@ -76,15 +76,17 @@ export class GPTModel extends tf.LayersModel { } setDebugLabel(label: string): void { - this.#debugLabel = label + this.#debugLabel = label; } setGoldfishLoss(config: GoldfishLossConfig | undefined): void { - this.#goldfishLoss = config?.enabled === true ? config : undefined + this.#goldfishLoss = config?.enabled === true ? config : undefined; } #debugMessage(message: string): string { - return this.#debugLabel === undefined ? message : `[${this.#debugLabel}] ${message}` + return this.#debugLabel === undefined + ? message + : `[${this.#debugLabel}] ${message}`; } override compile() { @@ -98,33 +100,40 @@ export class GPTModel extends tf.LayersModel { setLearningRate(lr: number): void { this.config.lr = lr; this.optimizer?.dispose(); - this.optimizer = this.config.weightDecay !== 0 - ? getCustomAdam(this, this.config.lr, this.config.weightDecay) - : tf.train.adam(this.config.lr); + this.optimizer = + this.config.weightDecay !== 0 + ? getCustomAdam(this, this.config.lr, this.config.weightDecay) + : tf.train.adam(this.config.lr); } - override async fitDataset(dataset: Dataset, trainingArgs: tf.ModelFitDatasetArgs & { iterationOffset?: number }): Promise { - const callbacks = trainingArgs.callbacks as tf.CustomCallbackArgs - const evalDataset = trainingArgs.validationData as tf.data.Dataset<{ xs: tf.Tensor2D, ys: tf.Tensor3D }> - const iterationOffset = trainingArgs.iterationOffset ?? 0 - await callbacks.onTrainBegin?.() + override async fitDataset( + dataset: Dataset, + trainingArgs: tf.ModelFitDatasetArgs & { iterationOffset?: number }, + ): Promise { + const callbacks = trainingArgs.callbacks as tf.CustomCallbackArgs; + const evalDataset = trainingArgs.validationData as tf.data.Dataset<{ + xs: tf.Tensor2D; + ys: tf.Tensor3D; + }>; + const iterationOffset = trainingArgs.iterationOffset ?? 0; + await callbacks.onTrainBegin?.(); for (let epoch = 1; epoch <= trainingArgs.epochs; epoch++) { let accuracyFraction: [number, number] = [0, 0]; - let averageLoss = 0 - let iteration = 1 + let averageLoss = 0; + let iteration = 1; - debug(this.#debugMessage("before iterator init")) - const iterator = await dataset.iterator() - debug(this.#debugMessage("after getting iterator, before next")) - let next = await iterator.next() - debug(this.#debugMessage("after next of iterator")) + debug(this.#debugMessage("before iterator init")); + const iterator = await dataset.iterator(); + debug(this.#debugMessage("after getting iterator, before next")); + let next = await iterator.next(); + debug(this.#debugMessage("after next of iterator")); while (next.done !== true && iteration <= this.config.maxIter) { - const reportedIteration = iterationOffset + iteration - let weightUpdateTime = performance.now() - await callbacks.onEpochBegin?.(epoch) - const { xs, ys } = next.value as { xs: tf.Tensor2D, ys: tf.Tensor3D } + const reportedIteration = iterationOffset + iteration; + let weightUpdateTime = performance.now(); + await callbacks.onEpochBegin?.(epoch); + const { xs, ys } = next.value as { xs: tf.Tensor2D; ys: tf.Tensor3D }; // TODO include as a tensor inside the model // const accTensor = tf.tidy(() => { @@ -145,44 +154,55 @@ export class GPTModel extends tf.LayersModel { // tf.dispose([accTensor]) accuracyFraction = [Number.NaN, Number.NaN]; - const goldfishLoss = this.#goldfishLoss + const goldfishLoss = this.#goldfishLoss; const goldfishMask = goldfishLoss === undefined ? undefined - : this.#buildGoldfishMask(xs, goldfishLoss) + : this.#buildGoldfishMask(xs, goldfishLoss); const lossTensor = tf.tidy(() => { - const { grads, value: lossTensor } = this.optimizer.computeGradients(() => { - const logits = this.apply(xs) - if (Array.isArray(logits)) - throw new Error('model outputs too many tensor') - if (logits instanceof tf.SymbolicTensor) - throw new Error('model outputs symbolic tensor') - return goldfishMask === undefined || goldfishLoss === undefined - ? tf.losses.softmaxCrossEntropy(ys, logits) - : this.#goldfishLossTensor(ys, logits, goldfishMask, goldfishLoss) - }) - const gradsClipped = clipByGlobalNormObj(grads, 1) - this.optimizer.applyGradients(gradsClipped) - tf.dispose(Object.values(gradsClipped)) - return lossTensor - }) - goldfishMask?.dispose() - - const loss = await lossTensor.array() - averageLoss += loss - weightUpdateTime = performance.now() - weightUpdateTime + const { grads, value: lossTensor } = this.optimizer.computeGradients( + () => { + const logits = this.apply(xs); + if (Array.isArray(logits)) + throw new Error("model outputs too many tensor"); + if (logits instanceof tf.SymbolicTensor) + throw new Error("model outputs symbolic tensor"); + return goldfishMask === undefined || goldfishLoss === undefined + ? tf.losses.softmaxCrossEntropy(ys, logits) + : this.#goldfishLossTensor( + ys, + logits, + goldfishMask, + goldfishLoss, + ); + }, + ); + const gradsClipped = clipByGlobalNormObj(grads, 1); + this.optimizer.applyGradients(gradsClipped); + tf.dispose(Object.values(gradsClipped)); + return lossTensor; + }); + goldfishMask?.dispose(); + + const loss = await lossTensor.array(); + averageLoss += loss; + weightUpdateTime = performance.now() - weightUpdateTime; if ( evalDataset !== undefined && this.config.evaluateEvery !== undefined && // iteration % this.config.evaluateEvery == 0 reportedIteration % this.config.evaluateEvery == 0 - ){ - const iterationLogs = await evaluate(this, evalDataset, this.config.maxEvalBatches) - debug(this.#debugMessage('evaluation metrics: %O'), iterationLogs); + ) { + const iterationLogs = await evaluate( + this, + evalDataset, + this.config.maxEvalBatches, + ); + debug(this.#debugMessage("evaluation metrics: %O"), iterationLogs); } - const memory = tf.memory().numBytes / 1024 / 1024 / 1024 + const memory = tf.memory().numBytes / 1024 / 1024 / 1024; debug(this.#debugMessage("training metrics: %O"), { epoch, iteration: reportedIteration, @@ -223,60 +243,61 @@ export class GPTModel extends tf.LayersModel { goldfishMask: tf.Tensor2D, config: GoldfishLossConfig, ): tf.Scalar { - if (Array.isArray(logits)) - throw new Error('model outputs too many tensor') - if (logits.rank !== 3) - throw new Error('model outputs wrong shape') + if (Array.isArray(logits)) throw new Error("model outputs too many tensor"); + if (logits.rank !== 3) throw new Error("model outputs wrong shape"); const tokenLosses = tf.neg( tf.sum(tf.mul(ys, tf.logSoftmax(logits as tf.Tensor3D, -1)), -1), - ) as tf.Tensor2D + ) as tf.Tensor2D; const supervisedMask = config.padTokenId === undefined ? goldfishMask : tf.mul( goldfishMask, - tf.cast(tf.notEqual(tf.argMax(ys, -1), config.padTokenId), 'float32'), - ) - - const denominator = tf.maximum(tf.sum(supervisedMask), tf.scalar(1)) - return tf.div(tf.sum(tf.mul(tokenLosses, supervisedMask)), denominator) + tf.cast( + tf.notEqual(tf.argMax(ys, -1), config.padTokenId), + "float32", + ), + ); + + const denominator = tf.maximum(tf.sum(supervisedMask), tf.scalar(1)); + return tf.div(tf.sum(tf.mul(tokenLosses, supervisedMask)), denominator); } #buildGoldfishMask( inputIds: tf.Tensor2D, config: GoldfishLossConfig, ): tf.Tensor2D { - const rows = inputIds.arraySync() + const rows = inputIds.arraySync(); const mask = rows.map((row) => row.map((_, targetOffset) => { - const targetIndex = targetOffset + 1 - const start = Math.max(0, targetIndex - config.h) - const context = row.slice(start, targetIndex) - return this.#hashTokenContext(context) % config.k === 0 ? 0 : 1 + const targetIndex = targetOffset + 1; + const start = Math.max(0, targetIndex - config.h); + const context = row.slice(start, targetIndex); + return this.#hashTokenContext(context) % config.k === 0 ? 0 : 1; }), - ) + ); - return tf.tensor2d(mask, inputIds.shape, 'float32') + return tf.tensor2d(mask, inputIds.shape, "float32"); } #hashTokenContext(tokens: number[]): number { - let hash = 0x811c9dc5 + let hash = 0x811c9dc5; for (const token of tokens) { - hash ^= token & 0xff - hash = Math.imul(hash, 0x01000193) - hash ^= (token >>> 8) & 0xff - hash = Math.imul(hash, 0x01000193) - hash ^= (token >>> 16) & 0xff - hash = Math.imul(hash, 0x01000193) - hash ^= (token >>> 24) & 0xff - hash = Math.imul(hash, 0x01000193) - hash ^= 0xff - hash = Math.imul(hash, 0x01000193) + hash ^= token & 0xff; + hash = Math.imul(hash, 0x01000193); + hash ^= (token >>> 8) & 0xff; + hash = Math.imul(hash, 0x01000193); + hash ^= (token >>> 16) & 0xff; + hash = Math.imul(hash, 0x01000193); + hash ^= (token >>> 24) & 0xff; + hash = Math.imul(hash, 0x01000193); + hash ^= 0xff; + hash = Math.imul(hash, 0x01000193); } - return hash >>> 0 + return hash >>> 0; } } diff --git a/discojs/src/models/index.ts b/discojs/src/models/index.ts index 2f2119dc7..56fd057c5 100644 --- a/discojs/src/models/index.ts +++ b/discojs/src/models/index.ts @@ -2,10 +2,10 @@ export { Model } from "./model.js"; export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js"; export { Tokenizer } from "./tokenizer.js"; -export { GPT } from './gpt/index.js' -export { ONNXModel } from './onnx.js' -export { GPTConfig } from './gpt/config.js' -export { GoldfishLossConfig } from './gpt/config.js' +export { GPT } from "./gpt/index.js"; +export { ONNXModel } from "./onnx.js"; +export { GPTConfig } from "./gpt/config.js"; +export { GoldfishLossConfig } from "./gpt/config.js"; export { evaluate as evaluate_hellaswag, HellaSwagDataset, diff --git a/discojs/src/privacy.ts b/discojs/src/privacy.ts index 54804cd84..823cafe69 100644 --- a/discojs/src/privacy.ts +++ b/discojs/src/privacy.ts @@ -6,11 +6,11 @@ import type { WeightNormHistory } from "./training/trainer.js"; /** Computes the Frobenius norm of the given weights. */ export async function frobeniusNorm(weights: tf.Tensor): Promise { - const squaredTensor = tf.tidy(() => weights.square().sum()); - const squared = await squaredTensor.data(); - squaredTensor.dispose(); - if (squared.length !== 1) throw new Error("unexpected weights shape"); - return Math.sqrt(squared[0]); + const squaredTensor = tf.tidy(() => weights.square().sum()); + const squared = await squaredTensor.data(); + squaredTensor.dispose(); + if (squared.length !== 1) throw new Error("unexpected weights shape"); + return Math.sqrt(squared[0]); } /** ALDP-FL implementation */ @@ -56,8 +56,8 @@ export async function addOptimalNoise( try { return clippedWeights.map((w, i) => - tf.tidy(() => w.add(tf.randomNormal(w.shape, 0, sigmas[i]))) - ) + tf.tidy(() => w.add(tf.randomNormal(w.shape, 0, sigmas[i]))), + ); } finally { clippedWeights.dispose(); } diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 4fd7dfc66..43e3c4476 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -7,7 +7,7 @@ import { GPTConfig } from "../models/index.js"; import * as coder from "./coder.js"; import { Encoded, isEncoded } from "./coder.js"; -import createDebug from "debug" +import createDebug from "debug"; const debug = createDebug("discojs:serialization:model"); @@ -35,9 +35,9 @@ export async function encode(model: Model): Promise { } export async function decode(encoded: Encoded): Promise> { - const raw = coder.decode(encoded) - - debug("IMPORTANT:model decoded") + const raw = coder.decode(encoded); + + debug("IMPORTANT:model decoded"); if (!Array.isArray(raw) || raw.length < 2) { throw new Error( @@ -45,16 +45,18 @@ export async function decode(encoded: Encoded): Promise> { ); } - debug("model encoding array length: %d", raw.length) + debug("model encoding array length: %d", raw.length); - const type = raw[0] as unknown - if (typeof type !== 'number') { - throw new Error('invalid encoding, first encoding field should be the model type') + const type = raw[0] as unknown; + if (typeof type !== "number") { + throw new Error( + "invalid encoding, first encoding field should be the model type", + ); } - debug("model type: %d", type) + debug("model type: %d", type); - const rawModel = raw[1] as unknown + const rawModel = raw[1] as unknown; switch (type) { case Type.TFJS: { debug("TFJS model decoding started"); @@ -87,8 +89,8 @@ export async function decode(encoded: Encoded): Promise> { if (raw.length == 2) { config = undefined; } else if (raw.length == 3) { - debug("GPT model config decoding") - config = raw[2] as GPTConfig + debug("GPT model config decoding"); + config = raw[2] as GPTConfig; } else { throw new Error( "invalid encoding, gpt-tfjs model encoding should be an array of length 2 or 3", @@ -100,12 +102,15 @@ export async function decode(encoded: Encoded): Promise> { "invalid encoding, gpt-tfjs model weights should be an encoding of its weights", ); - debug("GPT model weights decoding") - const weights = serialization.weights.decode(rawModel) + debug("GPT model weights decoding"); + const weights = serialization.weights.decode(rawModel); - debug("GPT model weights decoded") - debug("GPT model config: %O", config || "undefined, using default config") - return models.GPT.deserialize({weights, config}) + debug("GPT model weights decoded"); + debug( + "GPT model config: %O", + config || "undefined, using default config", + ); + return models.GPT.deserialize({ weights, config }); } default: throw new Error("invalid encoding, model type unrecognized"); diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index 568a7ca8e..6837a54f1 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -60,58 +60,58 @@ const nonLocalNetworkSchema = z ); export namespace TrainingInformation { - export const baseSchema = z.object({ - // number of epochs to run training for - epochs: z.number().positive().int(), - // number of epochs between each weight sharing round. - // e.g.if 3 then weights are shared every 3 epochs (in the distributed setting). - roundDuration: z.number().positive().int(), - // for GPT text tasks, number of training batches between each weight sharing round. - roundIterations: z.number().positive().int().optional(), - // run validation every N aggregation rounds. If 0, validation metrics are skipped. - validationFrequency: z.number().nonnegative().int().optional(), - // fraction of data to keep for validation, note this only works for image data - validationSplit: z.number().min(0).max(1), - // batch size of training data - batchSize: z.number().positive().int(), - // Tensor framework used by the model - tensorBackend: z.enum(["gpt", "tfjs"]), - }); + export const baseSchema = z.object({ + // number of epochs to run training for + epochs: z.number().positive().int(), + // number of epochs between each weight sharing round. + // e.g.if 3 then weights are shared every 3 epochs (in the distributed setting). + roundDuration: z.number().positive().int(), + // for GPT text tasks, number of training batches between each weight sharing round. + roundIterations: z.number().positive().int().optional(), + // run validation every N aggregation rounds. If 0, validation metrics are skipped. + validationFrequency: z.number().nonnegative().int().optional(), + // fraction of data to keep for validation, note this only works for image data + validationSplit: z.number().min(0).max(1), + // batch size of training data + batchSize: z.number().positive().int(), + // Tensor framework used by the model + tensorBackend: z.enum(["gpt", "tfjs"]), + }); - export const dataTypeToSchema = { - image: z.object({ - // classes, e.g. if two class of images, one with dogs and one with cats, then we would - // define ['dogs', 'cats']. - LABEL_LIST: z.array(z.string()).min(1), - // height of image to resize to - IMAGE_W: z.number().positive().int(), - // width of image to resize to - IMAGE_H: z.number().positive().int(), - }), - tabular: z.object({ - // the columns to be chosen as input data for the model - inputColumns: z.array(z.string()), - // the columns to be predicted by the model - outputColumn: z.string(), - }), - text: z.object({ - // should be set with the name of a Transformers.js pre-trained tokenizer, e.g., 'Xenova/gpt2'. - tokenizer: z.instanceof(Tokenizer), - // the maximum length of a input string used as input to a GPT model. It is used during preprocessing to - // truncate strings to a maximum length. The default value is tokenizer.model_max_length - contextLength: z.number().positive().int(), - // Goldfish loss drops a deterministic subset of shifted target-token losses while keeping full inputs. - goldfishLoss: z - .object({ - enabled: z.boolean(), - k: z.number().positive().int().default(4), - h: z.number().positive().int().default(13), - padTokenId: z.number().int().optional(), - }) - .optional(), - learningRate: z.number().positive().optional(), - }), - } satisfies Record; + export const dataTypeToSchema = { + image: z.object({ + // classes, e.g. if two class of images, one with dogs and one with cats, then we would + // define ['dogs', 'cats']. + LABEL_LIST: z.array(z.string()).min(1), + // height of image to resize to + IMAGE_W: z.number().positive().int(), + // width of image to resize to + IMAGE_H: z.number().positive().int(), + }), + tabular: z.object({ + // the columns to be chosen as input data for the model + inputColumns: z.array(z.string()), + // the columns to be predicted by the model + outputColumn: z.string(), + }), + text: z.object({ + // should be set with the name of a Transformers.js pre-trained tokenizer, e.g., 'Xenova/gpt2'. + tokenizer: z.instanceof(Tokenizer), + // the maximum length of a input string used as input to a GPT model. It is used during preprocessing to + // truncate strings to a maximum length. The default value is tokenizer.model_max_length + contextLength: z.number().positive().int(), + // Goldfish loss drops a deterministic subset of shifted target-token losses while keeping full inputs. + goldfishLoss: z + .object({ + enabled: z.boolean(), + k: z.number().positive().int().default(4), + h: z.number().positive().int().default(13), + padTokenId: z.number().int().optional(), + }) + .optional(), + learningRate: z.number().positive().optional(), + }), + } satisfies Record; export const networkToSchema = { decentralized: z diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 0236a4977..fe44cb80c 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -22,7 +22,7 @@ import { getAggregator } from "../aggregator/index.js"; import { enumerate, split } from "../utils/async_iterator.js"; import { EventEmitter } from "../utils/event_emitter.js"; -import createDebug from "debug" +import createDebug from "debug"; import { RoundLogs, Trainer } from "./trainer.js"; @@ -56,19 +56,19 @@ interface DiscoConfig { } export type SummaryLogs = { - round: number, - epoch: number, - trainingLoss: number, - trainingAccuracy: number, - peakMemory: number, - epochTime: number, - roundValidationLoss?: number, - roundValidationAccuracy?: number, - validationLoss?: number, - validationAccuracy?: number, - postAggregationValidationLoss?: number, - postAggregationValidationAccuracy?: number -} + round: number; + epoch: number; + trainingLoss: number; + trainingAccuracy: number; + peakMemory: number; + epochTime: number; + roundValidationLoss?: number; + roundValidationAccuracy?: number; + validationLoss?: number; + validationAccuracy?: number; + postAggregationValidationLoss?: number; + postAggregationValidationAccuracy?: number; +}; export type RoundStatus = | "not enough participants" // Server notification to wait for more participants @@ -83,19 +83,20 @@ function buildSummaryLog( epochLogs: EpochLogs, ): SummaryLogs { return { - round: roundNum, - epoch: epochNum, - trainingLoss: epochLogs.training.loss, - trainingAccuracy: epochLogs.training.accuracy, - peakMemory: epochLogs.peakMemory, - epochTime: epochLogs.epochTime, - roundValidationLoss: roundLogs.preRoundValidation?.loss, - roundValidationAccuracy: roundLogs.preRoundValidation?.accuracy, - validationLoss: epochLogs.validation?.loss, - validationAccuracy: epochLogs.validation?.accuracy, - postAggregationValidationLoss: roundLogs.postAggregationValidation?.loss, - postAggregationValidationAccuracy: roundLogs.postAggregationValidation?.accuracy, - } + round: roundNum, + epoch: epochNum, + trainingLoss: epochLogs.training.loss, + trainingAccuracy: epochLogs.training.accuracy, + peakMemory: epochLogs.peakMemory, + epochTime: epochLogs.epochTime, + roundValidationLoss: roundLogs.preRoundValidation?.loss, + roundValidationAccuracy: roundLogs.preRoundValidation?.accuracy, + validationLoss: epochLogs.validation?.loss, + validationAccuracy: epochLogs.validation?.accuracy, + postAggregationValidationLoss: roundLogs.postAggregationValidation?.loss, + postAggregationValidationAccuracy: + roundLogs.postAggregationValidation?.accuracy, + }; } /** @@ -205,13 +206,15 @@ export class Disco extends EventEmitter<{ dataset: Dataset, validationDataset?: Dataset, ): AsyncGenerator { - for await (const [roundNum, round] of enumerate(this.train(dataset, validationDataset))) { + for await (const [roundNum, round] of enumerate( + this.train(dataset, validationDataset), + )) { const [roundGen, roundLogsPromise] = async_iterator.split(round); const epochResults: Array<{ epochNum: number; epochLogs: EpochLogs }> = []; - debug("Starting round %d", roundNum) + debug("Starting round %d", roundNum); for await (const [epochNum, epoch] of enumerate(roundGen)) { const [epochGen, epochLogsPromise] = async_iterator.split(epoch); @@ -231,7 +234,10 @@ export class Disco extends EventEmitter<{ } /** Run whole train on dataset. */ - async trainFully(dataset: Dataset, validationDataset?: Dataset): Promise { + async trainFully( + dataset: Dataset, + validationDataset?: Dataset, + ): Promise { for await (const round of this.train(dataset, validationDataset)) for await (const epoch of round) for await (const _ of epoch); } @@ -263,7 +269,9 @@ export class Disco extends EventEmitter<{ this.#setModelTrainingOptions(this.trainer.model); debug("Initial model fetched successfully"); if (this.trainer.model === null) { - debug(`No pre-trained model provided for client, initializing randomly...`); + debug( + `No pre-trained model provided for client, initializing randomly...`, + ); } for await (const [roundNum, round] of enumerate( @@ -349,17 +357,23 @@ export class Disco extends EventEmitter<{ setLearningRate?: (learningRate: number) => void; }; - configurableModel.setGoldfishLoss?.(this.#task.trainingInformation.goldfishLoss); + configurableModel.setGoldfishLoss?.( + this.#task.trainingInformation.goldfishLoss, + ); if (this.#task.trainingInformation.goldfishLoss?.enabled === true) { const { k, h, padTokenId } = this.#task.trainingInformation.goldfishLoss; console.log( - `Using Goldfish loss with k=${k}, h=${h}` - + (padTokenId === undefined ? "" : `, padTokenId=${padTokenId}`), + `Using Goldfish loss with k=${k}, h=${h}` + + (padTokenId === undefined ? "" : `, padTokenId=${padTokenId}`), ); } if (this.#task.trainingInformation.learningRate !== undefined) { - configurableModel.setLearningRate?.(this.#task.trainingInformation.learningRate); - console.log(`Using GPT learning rate ${this.#task.trainingInformation.learningRate}`); + configurableModel.setLearningRate?.( + this.#task.trainingInformation.learningRate, + ); + console.log( + `Using GPT learning rate ${this.#task.trainingInformation.learningRate}`, + ); } } @@ -400,12 +414,22 @@ export class Disco extends EventEmitter<{ > { const { batchSize } = this.#task.trainingInformation; - let preprocessedTraining = processing.preprocess(this.#task, trainingDataset); - let preprocessedValidation = processing.preprocess(this.#task, validationDataset); + let preprocessedTraining = processing.preprocess( + this.#task, + trainingDataset, + ); + let preprocessedValidation = processing.preprocess( + this.#task, + validationDataset, + ); if (this.#preprocessOnce) { - preprocessedTraining = new Dataset(await arrayFromAsync(preprocessedTraining)); - preprocessedValidation = new Dataset(await arrayFromAsync(preprocessedValidation)); + preprocessedTraining = new Dataset( + await arrayFromAsync(preprocessedTraining), + ); + preprocessedValidation = new Dataset( + await arrayFromAsync(preprocessedValidation), + ); } return [ diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 98e93bebe..aeb140d95 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -40,7 +40,10 @@ type IterationTrainableTextModel = Model<"text"> & { ): AsyncGenerator; }; -function appendWeightHistory(weightNormHistory: WeightNormHistory, wc: number[]){ +function appendWeightHistory( + weightNormHistory: WeightNormHistory, + wc: number[], +) { return wc.reduce((hist, t, i) => { const arr = hist.get(i, List()); return hist.set(i, arr.push(t)); @@ -86,21 +89,36 @@ export class Trainer { this.#epochs = task.trainingInformation.epochs; this.#roundIterations = task.trainingInformation.roundIterations; this.#validationFrequency = task.trainingInformation.validationFrequency; - this.#shouldValidateAfterAggregation = task.trainingInformation.scheme !== "local"; - if ("privacy" in task.trainingInformation) - this.#privacy = task.trainingInformation.privacy; - - if (this.#roundIterations !== undefined && (task.dataType !== "text" || task.trainingInformation.tensorBackend !== "gpt")) + this.#shouldValidateAfterAggregation = + task.trainingInformation.scheme !== "local"; + if ("privacy" in task.trainingInformation) + this.#privacy = task.trainingInformation.privacy; + + if ( + this.#roundIterations !== undefined && + (task.dataType !== "text" || + task.trainingInformation.tensorBackend !== "gpt") + ) throw new Error("roundIterations is only supported for GPT text tasks"); - if (this.#roundIterations !== undefined && (!Number.isInteger(this.#roundIterations) || this.#roundIterations < 1)) + if ( + this.#roundIterations !== undefined && + (!Number.isInteger(this.#roundIterations) || this.#roundIterations < 1) + ) throw new Error("roundIterations must be a positive integer"); - if (this.#validationFrequency !== undefined && (!Number.isInteger(this.#validationFrequency) || this.#validationFrequency < 0)) + if ( + this.#validationFrequency !== undefined && + (!Number.isInteger(this.#validationFrequency) || + this.#validationFrequency < 0) + ) throw new Error("validationFrequency must be a non-negative integer"); // if (!Number.isInteger(this.#epochs / this.#roundDuration)) - if (this.#roundIterations === undefined && !Number.isInteger(this.#epochs / this.#roundDuration)) + if ( + this.#roundIterations === undefined && + !Number.isInteger(this.#epochs / this.#roundDuration) + ) throw new Error( `round duration ${this.#roundDuration} doesn't divide number of epochs ${this.#epochs}`, ); @@ -117,7 +135,7 @@ export class Trainer { AsyncGenerator, RoundLogs>, void > { - debug("Start train") + debug("Start train"); if (this.#training !== undefined) throw new Error( "training already running, stop it before launching a new one", @@ -144,7 +162,7 @@ export class Trainer { > { const totalRound = Math.trunc(this.#epochs / this.#roundDuration); - debug("Run rounds") + debug("Run rounds"); for (let round = 0; round < totalRound; round++) { await this.#client.onRoundBeginCommunication(); @@ -193,12 +211,17 @@ export class Trainer { while (pendingBatch !== undefined) { await this.#client.onRoundBeginCommunication(); - this.#previousRoundWeights = new WeightsContainer(this.model.weights.weights.map(t => t.clone())); + this.#previousRoundWeights = new WeightsContainer( + this.model.weights.weights.map((t) => t.clone()), + ); - let firstBatch: Batched | undefined = pendingBatch; + let firstBatch: Batched | undefined = + pendingBatch; pendingBatch = undefined; let done = false; - const prefixedIterator: AsyncIterator> = { + const prefixedIterator: AsyncIterator< + Batched + > = { next: async () => { if (firstBatch !== undefined) { const value = firstBatch; @@ -218,7 +241,7 @@ export class Trainer { prefixedIterator, this.#roundIterations, roundValidationDataset, - (roundDone) => done = roundDone, + (roundDone) => (done = roundDone), ); await this.#finishRoundCommunication(totalRound); @@ -237,7 +260,7 @@ export class Trainer { ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); - debug("Run round") + debug("Run round"); // Before starting the training, get the validation of global model const validation = @@ -269,9 +292,12 @@ export class Trainer { ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); - debug("Run iteration-based round") + debug("Run iteration-based round"); - const validation = validationDataset !== undefined ? await this.model.evaluate(validationDataset) : undefined; + const validation = + validationDataset !== undefined + ? await this.model.evaluate(validationDataset) + : undefined; const model = this.model as unknown as IterationTrainableTextModel; if (typeof model.trainNextBatches !== "function") @@ -279,9 +305,13 @@ export class Trainer { const [gen, result] = async_iterator.split( model.trainNextBatches( - datasetIterator as AsyncIterator>, + datasetIterator as AsyncIterator< + Batched + >, maxBatchCount, - validationDataset as Dataset> | undefined, + validationDataset as + | Dataset> + | undefined, setDone, ), ); @@ -311,7 +341,7 @@ export class Trainer { let disposeRoundWeightsAfterSend = false; try { - if (this.#privacy !== undefined){ + if (this.#privacy !== undefined) { if (this.#previousRoundWeights === undefined) throw new Error("previous round weights were not captured"); @@ -319,28 +349,33 @@ export class Trainer { const roundUpdate = roundWeights.sub(previousRoundWeights); try { const updateNorm = await Promise.all( - roundUpdate.weights.map(privacy.frobeniusNorm) + roundUpdate.weights.map(privacy.frobeniusNorm), + ); + this.#weightNormHistory = appendWeightHistory( + this.#weightNormHistory, + updateNorm, ); - this.#weightNormHistory = appendWeightHistory(this.#weightNormHistory, updateNorm); } finally { roundUpdate.dispose(); } const privateRoundWeights = await applyOptimalPrivacy( - previousRoundWeights, - roundWeights, - this.#privacy, - this.#weightNormHistory, - totalRound, - ); + previousRoundWeights, + roundWeights, + this.#privacy, + this.#weightNormHistory, + totalRound, + ); roundWeights = privateRoundWeights; disposeRoundWeightsAfterSend = true; } - const networkWeights = await this.#client.onRoundEndCommunication(roundWeights); + const networkWeights = + await this.#client.onRoundEndCommunication(roundWeights); this.model.weights = networkWeights; - return this.#shouldValidateAfterAggregation && validationDataset !== undefined + return this.#shouldValidateAfterAggregation && + validationDataset !== undefined ? await this.model.evaluate(validationDataset) : undefined; } finally { @@ -367,26 +402,26 @@ async function applyOptimalPrivacy( ): Promise { let ret = current; - // Clipping radius for BFT - if ("byzantineFaultTolerance" in options) { - // might need to change the variable name - const previousRoundWeights = - previous ?? current.map((w) => tf.zerosLike(w)); - const weightsProgress = current.sub(previousRoundWeights); - const clippedProgress = await privacy.clipNorm( - weightsProgress, - Repeat(options.byzantineFaultTolerance.clippingRadius) - .take(weightsProgress.weights.length) - .toArray(), - ); - try { - ret = previousRoundWeights.add(clippedProgress); - } finally { - weightsProgress.dispose(); - clippedProgress.dispose(); - if (previous === undefined) previousRoundWeights.dispose(); - } - } + // Clipping radius for BFT + if ("byzantineFaultTolerance" in options) { + // might need to change the variable name + const previousRoundWeights = + previous ?? current.map((w) => tf.zerosLike(w)); + const weightsProgress = current.sub(previousRoundWeights); + const clippedProgress = await privacy.clipNorm( + weightsProgress, + Repeat(options.byzantineFaultTolerance.clippingRadius) + .take(weightsProgress.weights.length) + .toArray(), + ); + try { + ret = previousRoundWeights.add(clippedProgress); + } finally { + weightsProgress.dispose(); + clippedProgress.dispose(); + if (previous === undefined) previousRoundWeights.dispose(); + } + } // Adding Gaussian noise for DP const dpOptions = options.differentialPrivacy; @@ -414,32 +449,32 @@ async function applyOptimalPrivacy( ) : dpClippingRadius; - const sigmas = effectiveRadius.map((r) => - (2 * r * Math.sqrt(2 * Math.log(1.25 / delta))) / epsilon, - ); - debug("DP applied: %O", { - totalRound, - epsilon, - delta, - radiusMin: Math.min(...effectiveRadius), - radiusMax: Math.max(...effectiveRadius), - sigmaMin: Math.min(...sigmas), - sigmaMax: Math.max(...sigmas), - }); - - const noisyProgress = await privacy.addOptimalNoise( - weightsProgress, - epsilon, - delta, - effectiveRadius, - ); - try { - ret = previousEpochWeights.add(noisyProgress); - } finally { - weightsProgress.dispose(); - noisyProgress.dispose(); - if (previous === undefined) previousEpochWeights.dispose(); - } - } - return ret; + const sigmas = effectiveRadius.map( + (r) => (2 * r * Math.sqrt(2 * Math.log(1.25 / delta))) / epsilon, + ); + debug("DP applied: %O", { + totalRound, + epsilon, + delta, + radiusMin: Math.min(...effectiveRadius), + radiusMax: Math.max(...effectiveRadius), + sigmaMin: Math.min(...sigmas), + sigmaMax: Math.max(...sigmas), + }); + + const noisyProgress = await privacy.addOptimalNoise( + weightsProgress, + epsilon, + delta, + effectiveRadius, + ); + try { + ret = previousEpochWeights.add(noisyProgress); + } finally { + weightsProgress.dispose(); + noisyProgress.dispose(); + if (previous === undefined) previousEpochWeights.dispose(); + } + } + return ret; } diff --git a/onnx-converter/src/convert_onnx.ts b/onnx-converter/src/convert_onnx.ts index 9e8422d2d..43767d7cd 100644 --- a/onnx-converter/src/convert_onnx.ts +++ b/onnx-converter/src/convert_onnx.ts @@ -8,8 +8,8 @@ import { models, serialization } from "@epfml/discojs"; const OUTPUT_FILENAME = "model.json"; const GPT2_N_LAYER = 12; const GPT2_CONTEXT_LENGTH = 1024; -const ONNX_URL = "https://huggingface.co/Xenova/gpt2/resolve/main/onnx/decoder_model.onnx?download=true" - +const ONNX_URL = + "https://huggingface.co/Xenova/gpt2/resolve/main/onnx/decoder_model.onnx?download=true"; async function main() { console.log(`Downloading ONNX model from ${ONNX_URL}...`); @@ -33,7 +33,10 @@ async function main() { console.log("ONNX model loaded successfully"); // Init empty TF.js model - const gptModel = new models.GPT({ modelType: 'gpt2', contextLength: GPT2_CONTEXT_LENGTH }); + const gptModel = new models.GPT({ + modelType: "gpt2", + contextLength: GPT2_CONTEXT_LENGTH, + }); if (gptModel.config.nLayer != GPT2_N_LAYER) throw new Error( `ONNX conversion only supports GPT-2 with 12 layers, instead found ${gptModel.config.nLayer}.`, @@ -59,12 +62,16 @@ async function main() { throw new Error(`Undefined layer dimensions for ${tensor.name}`); const dims = tensor.dims.map((d) => Number(d)); const flatData = parseTensorData(tensor); - let tfTensor = tf.tensor(flatData).reshape(dims) + let tfTensor = tf.tensor(flatData).reshape(dims); if (tensor.name === "transformer.wpe.weight") { if (dims.length !== 2) - throw new Error(`Expected transformer.wpe.weight to be a 2D tensor, got ${dims.length}D.`); + throw new Error( + `Expected transformer.wpe.weight to be a 2D tensor, got ${dims.length}D.`, + ); if (dims[0] < GPT2_CONTEXT_LENGTH) - throw new Error(`ONNX positional embeddings only support context length ${dims[0]}, requested ${GPT2_CONTEXT_LENGTH}.`); + throw new Error( + `ONNX positional embeddings only support context length ${dims[0]}, requested ${GPT2_CONTEXT_LENGTH}.`, + ); tfTensor = tfTensor.slice([0, 0], [GPT2_CONTEXT_LENGTH, dims[1]]); } preTrainedWeights = preTrainedWeights.set(tfjsName, tfTensor); diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 4c8e94337..133650f18 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -18,25 +18,25 @@ import FederatedMessages = client.federated.messages; const debug = createDebug("server:controllers:federated"); function debugProcessMemory(label: string): void { - const m = process.memoryUsage() + const m = process.memoryUsage(); debug("%s memory: %O", label, { rssGB: m.rss / 1024 / 1024 / 1024, heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, externalGB: m.external / 1024 / 1024 / 1024, arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, - }) + }); } export class FederatedController extends TrainingController< D, "federated" > { - #pendingUpdateRecipients = new Map() + #pendingUpdateRecipients = new Map(); /** * Aggregators for each hosted task. By default the server waits for 100% of the nodes to send their contributions before aggregating the updates */ - #aggregator = this.#makeAggregator() + #aggregator = this.#makeAggregator(); /** * The most up to date global weights. The model weights are already serialized and * can be sent to participants, before starting training, or when joining mid-training @@ -44,58 +44,72 @@ export class FederatedController extends TrainingController< */ #latestGlobalWeights: serialization.Encoded; - constructor( - task: Task, - private readonly initialWeights: serialization.Encoded, - ) { - super(task) - this.#latestGlobalWeights = this.initialWeights + constructor( + task: Task, + private readonly initialWeights: serialization.Encoded, + ) { + super(task); + this.#latestGlobalWeights = this.initialWeights; } #makeAggregator(): aggregators.MeanAggregator { - const aggregator = new aggregators.MeanAggregator(undefined, 1, 'relative') + const aggregator = new aggregators.MeanAggregator(undefined, 1, "relative"); - aggregator.on('aggregation', async (weightUpdate) => { + aggregator.on("aggregation", async (weightUpdate) => { try { - const recipients = this.#pendingUpdateRecipients - this.#pendingUpdateRecipients = new Map() + const recipients = this.#pendingUpdateRecipients; + this.#pendingUpdateRecipients = new Map(); - debugProcessMemory(`round ${aggregator.round} before encoding aggregate`) - const payload = await serialization.weights.encode(weightUpdate) - debugProcessMemory(`round ${aggregator.round} after encoding aggregate`) - debug("round %o aggregate payload byteLength=%d", aggregator.round, payload.byteLength) - this.#latestGlobalWeights = payload + debugProcessMemory( + `round ${aggregator.round} before encoding aggregate`, + ); + const payload = await serialization.weights.encode(weightUpdate); + debugProcessMemory( + `round ${aggregator.round} after encoding aggregate`, + ); + debug( + "round %o aggregate payload byteLength=%d", + aggregator.round, + payload.byteLength, + ); + this.#latestGlobalWeights = payload; const msg: FederatedMessages.ReceiveServerPayload = { type: MessageTypes.ReceiveServerPayload, round: aggregator.round, payload, - nbOfParticipants: this.connections.size - } - const encodedMsg = msgpack.encode(msg) - debugProcessMemory(`round ${aggregator.round} after encoding websocket message`) + nbOfParticipants: this.connections.size, + }; + const encodedMsg = msgpack.encode(msg); + debugProcessMemory( + `round ${aggregator.round} after encoding websocket message`, + ); recipients.forEach((recipientWs, recipientId) => { debug( "Sending global weights for round %o to client [%s]", aggregator.round, recipientId.slice(0, 4), - ) - recipientWs.send(encodedMsg) + ); + recipientWs.send(encodedMsg); debug( "Aggregated payload sent to client [%s] for round %o", recipientId.slice(0, 4), aggregator.round, - ) - }) - debugProcessMemory(`round ${aggregator.round} after sending aggregate to ${recipients.size} clients`) + ); + }); + debugProcessMemory( + `round ${aggregator.round} after sending aggregate to ${recipients.size} clients`, + ); } finally { - weightUpdate.dispose() - debugProcessMemory(`round ${aggregator.round} after disposing aggregate tensor`) + weightUpdate.dispose(); + debugProcessMemory( + `round ${aggregator.round} after disposing aggregate tensor`, + ); } - }) + }); - return aggregator + return aggregator; } /** @@ -119,9 +133,9 @@ export class FederatedController extends TrainingController< } const shortId = clientId.slice(0, 4); - ws.on('error', (err) => { - debug("websocket error for client [%s]: %o", shortId, err) - }) + ws.on("error", (err) => { + debug("websocket error for client [%s]: %o", shortId, err); + }); // Setup callbacks triggered upon receiving the different client messages ws.on("message", (data: Buffer) => { @@ -145,8 +159,12 @@ export class FederatedController extends TrainingController< const msg: FederatedMessages.NewFederatedNodeInfo = { type: MessageTypes.NewFederatedNodeInfo, id: clientId, - waitForMoreParticipants: this.connections.size < minNbOfParticipants, - payload: this.#aggregator.round === 0 ? undefined : this.#latestGlobalWeights, + waitForMoreParticipants: + this.connections.size < minNbOfParticipants, + payload: + this.#aggregator.round === 0 + ? undefined + : this.#latestGlobalWeights, round: this.#aggregator.round, nbOfParticipants: this.connections.size, }; @@ -166,19 +184,26 @@ export class FederatedController extends TrainingController< shortId, round, this.connections.size, - ) - const weights = serialization.weights.decode(payload) - let added = false + ); + const weights = serialization.weights.decode(payload); + let added = false; try { // Add the contribution - debug("Adding contribution from client [%s] to aggregator for round %d", shortId, round) - this.#pendingUpdateRecipients.set(clientId, ws) - this.#aggregator.add(clientId, weights, round) - added = true - debug(`Successfully added contribution from client [%s] for round ${round}`, shortId) + debug( + "Adding contribution from client [%s] to aggregator for round %d", + shortId, + round, + ); + this.#pendingUpdateRecipients.set(clientId, ws); + this.#aggregator.add(clientId, weights, round); + added = true; + debug( + `Successfully added contribution from client [%s] for round ${round}`, + shortId, + ); } finally { - weights.dispose() - if (!added) this.#pendingUpdateRecipients.delete(clientId) + weights.dispose(); + if (!added) this.#pendingUpdateRecipients.delete(clientId); } } else { // If the client sent an invalid or outdated contribution @@ -207,17 +232,17 @@ export class FederatedController extends TrainingController< // Setup callback for client leaving the session ws.on("close", () => { // Remove the participant when the websocket is closed - this.connections = this.connections.delete(clientId) - this.#pendingUpdateRecipients.delete(clientId) - this.#aggregator.removeNode(clientId) - debug("client [%s] left", shortId) + this.connections = this.connections.delete(clientId); + this.#pendingUpdateRecipients.delete(clientId); + this.#aggregator.removeNode(clientId); + debug("client [%s] left", shortId); // Reset the training session when all participants left if (this.connections.size === 0) { - debug("All participants left. Resetting the training session") - this.#pendingUpdateRecipients.clear() - this.#aggregator = this.#makeAggregator() - this.#latestGlobalWeights = this.initialWeights + debug("All participants left. Resetting the training session"); + this.#pendingUpdateRecipients.clear(); + this.#aggregator = this.#makeAggregator(); + this.#latestGlobalWeights = this.initialWeights; } // Check if we dropped below the minimum number of participant required diff --git a/server/src/routes/training_router.ts b/server/src/routes/training_router.ts index b0186fd76..8c1ff3560 100644 --- a/server/src/routes/training_router.ts +++ b/server/src/routes/training_router.ts @@ -55,15 +55,15 @@ export class TrainingRouter> { // The federated controller takes the initial model weights at initialization // so that it can send it to new clients - const model = await serialization.model.decode(encodedModel) - const weights = model.weights - let encodedWeights: serialization.Encoded + const model = await serialization.model.decode(encodedModel); + const weights = model.weights; + let encodedWeights: serialization.Encoded; try { - encodedWeights = await serialization.weights.encode(weights) + encodedWeights = await serialization.weights.encode(weights); } finally { - model[Symbol.dispose]() + model[Symbol.dispose](); } - taskController = new FederatedController(t, encodedWeights) + taskController = new FederatedController(t, encodedWeights); } else { const t = task as Task; diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts index 2561fe79f..563d8de10 100644 --- a/server/tests/e2e/federated.spec.ts +++ b/server/tests/e2e/federated.spec.ts @@ -151,24 +151,26 @@ describe("end-to-end federated", () => { assert.isTrue(m1.equals(m2)); }); - it("two wikitext reach consensus", { timeout: 500_000 }, async () => { - const task = await defaultTasks.wikitext.getTask(); - task.trainingInformation = { - ...task.trainingInformation, - epochs: 2, - roundDuration: 2, - minNbOfParticipants: 2, - }; - const url = await startServer({ - ...defaultTasks.wikitext, - getModel: () => - Promise.resolve(new models.GPT({ - contextLength: task.trainingInformation.contextLength, - maxIter: 10, - })), - getTask: () => Promise.resolve(task), - }); - const dataset = datasets.loadWikitext(); + it("two wikitext reach consensus", { timeout: 500_000 }, async () => { + const task = await defaultTasks.wikitext.getTask(); + task.trainingInformation = { + ...task.trainingInformation, + epochs: 2, + roundDuration: 2, + minNbOfParticipants: 2, + }; + const url = await startServer({ + ...defaultTasks.wikitext, + getModel: () => + Promise.resolve( + new models.GPT({ + contextLength: task.trainingInformation.contextLength, + maxIter: 10, + }), + ), + getTask: () => Promise.resolve(task), + }); + const dataset = datasets.loadWikitext(); const [r1, r2] = await Promise.all([ runUser(url, task, dataset, false), From 39f9689f290982760c35f6bc3b7d6526e15da33e Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 8 Jul 2026 16:05:06 +0200 Subject: [PATCH 56/89] fix lint --- .knip.json | 5 ++++- cli/package.json | 1 + cli/src/args.ts | 15 +++------------ discojs/src/models/gpt/model.ts | 2 +- discojs/src/training/disco.ts | 2 +- pnpm-lock.yaml | 3 +++ 6 files changed, 13 insertions(+), 15 deletions(-) diff --git a/.knip.json b/.knip.json index f33580dd0..96ebcc884 100644 --- a/.knip.json +++ b/.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": { diff --git a/cli/package.json b/cli/package.json index 175d204f7..f487f7579 100644 --- a/cli/package.json +++ b/cli/package.json @@ -22,6 +22,7 @@ "@epfml/discojs": "workspace:", "@epfml/discojs-node": "workspace:", "@tensorflow/tfjs-node": "catalog:", + "debug": "catalog:", "immutable": "catalog:", "server": "workspace:" }, diff --git a/cli/src/args.ts b/cli/src/args.ts index 43ecda229..6841f9760 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -287,18 +287,9 @@ export const args: BenchmarkArguments = { task.trainingInformation.roundDuration = unsafeArgs.roundDuration; task.trainingInformation.epochs = unsafeArgs.epochs; task.trainingInformation.validationSplit = unsafeArgs.validationSplit; - ( - task.trainingInformation as typeof task.trainingInformation & { - roundIterations?: number; - validationFrequency?: number; - } - ).roundIterations = unsafeArgs.roundIterations; - ( - task.trainingInformation as typeof task.trainingInformation & { - roundIterations?: number; - validationFrequency?: number; - } - ).validationFrequency = unsafeArgs.validationFrequency; + task.trainingInformation.roundIterations = unsafeArgs.roundIterations; + task.trainingInformation.validationFrequency = + unsafeArgs.validationFrequency; if (unsafeArgs.goldfishLoss) { if ( diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index cc9506537..5e877c21b 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -248,7 +248,7 @@ export class GPTModel extends tf.LayersModel { const tokenLosses = tf.neg( tf.sum(tf.mul(ys, tf.logSoftmax(logits as tf.Tensor3D, -1)), -1), - ) as tf.Tensor2D; + ); const supervisedMask = config.padTokenId === undefined diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index fe44cb80c..5926bbe99 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -264,7 +264,7 @@ export class Disco extends EventEmitter<{ // the client fetches the latest weights upon connection // TODO unsafe cast debug("Connecting to client and fetching initial model..."); - this.trainer.model = (await this.#client.connect()) as Model; + this.trainer.model = await this.#client.connect(); this.#setModelDebugLabel(this.trainer.model); this.#setModelTrainingOptions(this.trainer.model); debug("Initial model fetched successfully"); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index f62d4dea1..5a8f27eb2 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -289,6 +289,9 @@ importers: '@tensorflow/tfjs-node': specifier: 'catalog:' version: 4.22.0(seedrandom@3.0.5)(supports-color@8.1.1) + debug: + specifier: 'catalog:' + version: 4.4.3(supports-color@8.1.1) immutable: specifier: 'catalog:' version: 5.1.6 From a9dfeef340f7bf2f74cec1912b22a13016fd5dfe Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 8 Jul 2026 16:09:26 +0200 Subject: [PATCH 57/89] fix formatting --- cli/package.json | 1 + 1 file changed, 1 insertion(+) diff --git a/cli/package.json b/cli/package.json index f487f7579..647d2c461 100644 --- a/cli/package.json +++ b/cli/package.json @@ -21,6 +21,7 @@ "dependencies": { "@epfml/discojs": "workspace:", "@epfml/discojs-node": "workspace:", + "@tensorflow/tfjs": "catalog:", "@tensorflow/tfjs-node": "catalog:", "debug": "catalog:", "immutable": "catalog:", From 3446d0d34493fa15980ad9682f482d55d18fa876 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 8 Jul 2026 16:11:46 +0200 Subject: [PATCH 58/89] fix formatting --- cli/package.json | 2 +- pnpm-lock.yaml | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/cli/package.json b/cli/package.json index 647d2c461..b517de895 100644 --- a/cli/package.json +++ b/cli/package.json @@ -21,7 +21,7 @@ "dependencies": { "@epfml/discojs": "workspace:", "@epfml/discojs-node": "workspace:", - "@tensorflow/tfjs": "catalog:", + "@tensorflow/tfjs": "catalog:", "@tensorflow/tfjs-node": "catalog:", "debug": "catalog:", "immutable": "catalog:", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 5a8f27eb2..5ce0e47a3 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -286,6 +286,9 @@ importers: '@epfml/discojs-node': specifier: 'workspace:' version: link:../discojs-node + '@tensorflow/tfjs': + specifier: 'catalog:' + version: 4.22.0(seedrandom@3.0.5) '@tensorflow/tfjs-node': specifier: 'catalog:' version: 4.22.0(seedrandom@3.0.5)(supports-color@8.1.1) From 34e16c89cc745183caf23f4e3cf077839d845abd Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Wed, 8 Jul 2026 16:22:32 +0200 Subject: [PATCH 59/89] add part --- discojs/src/models/gpt/model.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 5e877c21b..acbd03d15 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -135,6 +135,10 @@ export class GPTModel extends tf.LayersModel { await callbacks.onEpochBegin?.(epoch); const { xs, ys } = next.value as { xs: tf.Tensor2D; ys: tf.Tensor3D }; + let preprocessingTime = performance.now(); + await Promise.all([xs.data(), ys.data()]); + preprocessingTime = performance.now() - preprocessingTime; + // TODO include as a tensor inside the model // const accTensor = tf.tidy(() => { // const logits = this.apply(xs) From 434a01885c241e9a3cfe3de8113fddad0e111baf Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Jul 2026 13:28:21 +0200 Subject: [PATCH 60/89] Update discojs/src/training/trainer.ts Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com> --- discojs/src/training/trainer.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index aeb140d95..1e5fb2b48 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -114,7 +114,6 @@ export class Trainer { ) throw new Error("validationFrequency must be a non-negative integer"); - // if (!Number.isInteger(this.#epochs / this.#roundDuration)) if ( this.#roundIterations === undefined && !Number.isInteger(this.#epochs / this.#roundDuration) From b861f555db7aa02507e117e11988b74a7c73d851 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Jul 2026 13:28:58 +0200 Subject: [PATCH 61/89] Update discojs/src/training/trainer.ts Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com> --- discojs/src/training/trainer.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 1e5fb2b48..d6fe6e29d 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -141,7 +141,6 @@ export class Trainer { ); try { - // this.#training = this.#runRounds(dataset, validationDataset); this.#training = this.#roundIterations === undefined ? this.#runRounds(dataset, validationDataset) From f1d486b920c417d9b887bc278a8afe154b50b3cb Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Jul 2026 13:29:55 +0200 Subject: [PATCH 62/89] Update discojs/src/training/trainer.ts Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com> --- discojs/src/training/trainer.ts | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index d6fe6e29d..9fc663cb2 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -301,7 +301,7 @@ export class Trainer { if (typeof model.trainNextBatches !== "function") throw new Error("model does not support iteration-based training"); - const [gen, result] = async_iterator.split( + const [gen, epochLogs] = async_iterator.split( model.trainNextBatches( datasetIterator as AsyncIterator< Batched @@ -315,8 +315,7 @@ export class Trainer { ); yield gen; - const epochLogs = await result; - epochsLogs = epochsLogs.push(epochLogs); + epochsLogs = epochsLogs.push(await epochLogs); return { epochs: epochsLogs, From c34faced60e355c5869f5fd4556fdb8967ff1fec Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Jul 2026 13:30:21 +0200 Subject: [PATCH 63/89] Update discojs/src/training/trainer.ts Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com> --- discojs/src/training/trainer.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 9fc663cb2..29baa7f6f 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -203,10 +203,7 @@ export class Trainer { for (let epoch = 0; epoch < this.#epochs; epoch++) { const trainingIterator = dataset[Symbol.asyncIterator](); let next = await trainingIterator.next(); - let pendingBatch: Batched | undefined = - next.done === true ? undefined : next.value; - - while (pendingBatch !== undefined) { + while (next.done !== true) { await this.#client.onRoundBeginCommunication(); this.#previousRoundWeights = new WeightsContainer( From 420787777865432a300f97215db91cb5d7d6b559 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Jul 2026 13:45:06 +0200 Subject: [PATCH 64/89] clean debug logs --- datasets/.gitignore | 1 - discojs/src/serialization/model.ts | 14 -------------- 2 files changed, 15 deletions(-) diff --git a/datasets/.gitignore b/datasets/.gitignore index bdbf3ed6d..e644c626a 100644 --- a/datasets/.gitignore +++ b/datasets/.gitignore @@ -7,7 +7,6 @@ /simple_face-example.png /titanic* /mnist* -/medicalMCQtxtNoExplanation # wikitext /wikitext/ diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 43e3c4476..e5277b404 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -20,13 +20,11 @@ export async function encode(model: Model): Promise { switch (true) { case model instanceof models.TFJS: { const serialized = await model.serialize(); - debug("TFJS model serialized"); return coder.encode([Type.TFJS, ...serialized]); } case model instanceof models.GPT: { const { weights, config } = model.serialize(); const serializedWeights = await serialization.weights.encode(weights); - debug("GPT model weights serialized"); return coder.encode([Type.GPT, serializedWeights, config]); } default: @@ -37,16 +35,12 @@ export async function encode(model: Model): Promise { export async function decode(encoded: Encoded): Promise> { const raw = coder.decode(encoded); - debug("IMPORTANT:model decoded"); - if (!Array.isArray(raw) || raw.length < 2) { throw new Error( "invalid encoding, encoding isn't an array or doesn't contain enough values", ); } - debug("model encoding array length: %d", raw.length); - const type = raw[0] as unknown; if (typeof type !== "number") { throw new Error( @@ -54,20 +48,15 @@ export async function decode(encoded: Encoded): Promise> { ); } - debug("model type: %d", type); - const rawModel = raw[1] as unknown; switch (type) { case Type.TFJS: { - debug("TFJS model decoding started"); if (raw.length !== 3) throw new Error( "invalid TFJS model encoding: should be an array of length 3", ); const [rawDatatype, rawModel] = raw.slice(1) as unknown[]; - debug("TFJS model datatype: %s", rawDatatype); - let datatype; switch (rawDatatype) { case "image": @@ -89,7 +78,6 @@ export async function decode(encoded: Encoded): Promise> { if (raw.length == 2) { config = undefined; } else if (raw.length == 3) { - debug("GPT model config decoding"); config = raw[2] as GPTConfig; } else { throw new Error( @@ -102,10 +90,8 @@ export async function decode(encoded: Encoded): Promise> { "invalid encoding, gpt-tfjs model weights should be an encoding of its weights", ); - debug("GPT model weights decoding"); const weights = serialization.weights.decode(rawModel); - debug("GPT model weights decoded"); debug( "GPT model config: %O", config || "undefined, using default config", From ff621151075928f6ea3c7141ee59bb212483e336 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Sun, 19 Jul 2026 14:56:01 +0200 Subject: [PATCH 65/89] add hash comment and fix refactor --- discojs/src/models/gpt/model.ts | 6 ++++++ discojs/src/training/trainer.ts | 4 +--- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index acbd03d15..379a3f25d 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -286,6 +286,12 @@ export class GPTModel extends tf.LayersModel { return tf.tensor2d(mask, inputIds.shape, "float32"); } + /** + * Computes a deterministic 32-bit FNV-1a-style hash of a token context. + * Each token ID is mixed in little-endian byte order, followed by a separator + * byte to preserve token boundaries. Goldfish uses the unsigned result modulo + * `k` to consistently select which target-token losses to drop. + */ #hashTokenContext(tokens: number[]): number { let hash = 0x811c9dc5; diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 29baa7f6f..483675b8f 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -211,8 +211,7 @@ export class Trainer { ); let firstBatch: Batched | undefined = - pendingBatch; - pendingBatch = undefined; + next.value; let done = false; const prefixedIterator: AsyncIterator< Batched @@ -244,7 +243,6 @@ export class Trainer { round++; if (done) break; next = await trainingIterator.next(); - pendingBatch = next.done === true ? undefined : next.value; } } } From 3a921b90451bcb5f9e455a95b4f0151dba08f30a Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 27 Aug 2026 17:29:12 +0200 Subject: [PATCH 66/89] make changes after first part of code review --- cli/src/args.ts | 24 +++- .../src/client/federated/federated_client.ts | 30 ----- discojs/src/models/gpt/model.ts | 2 + discojs/src/task/training_information.ts | 2 + discojs/src/training/disco.ts | 14 +-- discojs/src/training/trainer.ts | 117 +++++++++++++----- .../src/controllers/federated_controller.ts | 66 +++++----- server/src/routes/training_router.ts | 6 + 8 files changed, 150 insertions(+), 111 deletions(-) diff --git a/cli/src/args.ts b/cli/src/args.ts index 6841f9760..61d326ac6 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -19,9 +19,25 @@ export interface BenchmarkArguments { 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; 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; @@ -101,13 +117,14 @@ const unsafeArgs = parse( 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, }, validationFrequency: { type: Number, description: - "Run validation every N aggregation rounds. Defaults to every round; use 0 to disable validation metrics.", + "Validate the first aggregation round and every N rounds after it. Defaults to every round; use 0 to disable validation metrics.", optional: true, }, datasetPath: { @@ -119,7 +136,8 @@ const unsafeArgs = parse( validationDatasetPath: { type: String, alias: "V", - description: "Path to the validation dataset", + description: + "Path to a separate validation dataset shared by all clients. Takes precedence over --validationSplit.", optional: true, }, outputPath: { diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 1f5fbddb7..42b3fea15 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -9,18 +9,6 @@ import * as messages from "./messages.js"; const debug = createDebug("discojs:client:federated"); -function debugProcessMemory(label: string): void { - if (typeof process === "undefined") return; - - const m = process.memoryUsage(); - debug("%s memory: %O", label, { - rssGB: m.rss / 1024 / 1024 / 1024, - heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, - externalGB: m.external / 1024 / 1024 / 1024, - arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, - }); -} - /** * Arbitrary node id assigned to the federated server which we are communicating with. * Indeed, the server acts as a node within the network. In the federated setting described @@ -156,13 +144,7 @@ export class FederatedClient extends Client<"federated"> { const round = this.aggregator.round; { - debugProcessMemory( - `[${shortenId(this.ownId)}] round ${round} before encode`, - ); const payload = await serialization.weights.encode(payloadToServer); - debugProcessMemory( - `[${shortenId(this.ownId)}] round ${round} after encode`, - ); debug( "[%s] encoded payload for round %d byteLength=%d", shortenId(this.ownId), @@ -178,13 +160,7 @@ export class FederatedClient extends Client<"federated"> { // Need to await the resulting global model right after sending our local contribution // to make sure we don't miss it - debugProcessMemory( - `[${shortenId(this.ownId)}] round ${round} before send`, - ); this.server.send(msg); - debugProcessMemory( - `[${shortenId(this.ownId)}] round ${round} after send`, - ); } debug( `[${shortenId(this.ownId)}] sent its local update to the server for round ${round}`, @@ -198,13 +174,7 @@ export class FederatedClient extends Client<"federated"> { nbOfParticipants, } = await waitMessage(this.server, type.ReceiveServerPayload); // Wait indefinitely for the server update this.nbOfParticipants = nbOfParticipants; // Save the current participants - debugProcessMemory( - `[${shortenId(this.ownId)}] round ${round} after receive`, - ); const serverResult = serialization.weights.decode(payloadFromServer); - debugProcessMemory( - `[${shortenId(this.ownId)}] round ${round} after decode server payload`, - ); this.aggregator.setRound(serverRound); return serverResult; diff --git a/discojs/src/models/gpt/model.ts b/discojs/src/models/gpt/model.ts index 379a3f25d..892dd3aaf 100644 --- a/discojs/src/models/gpt/model.ts +++ b/discojs/src/models/gpt/model.ts @@ -190,6 +190,8 @@ export class GPTModel extends tf.LayersModel { goldfishMask?.dispose(); const loss = await lossTensor.array(); + lossTensor.dispose(); + tf.dispose([xs, ys]); averageLoss += loss; weightUpdateTime = performance.now() - weightUpdateTime; diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index 6837a54f1..8251d9872 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -70,6 +70,8 @@ export namespace TrainingInformation { roundIterations: z.number().positive().int().optional(), // run validation every N aggregation rounds. If 0, validation metrics are skipped. validationFrequency: z.number().nonnegative().int().optional(), + // whether to validate before aggregation, after aggregation, or at both points + validationMode: z.enum(["before", "after", "both"]).optional(), // fraction of data to keep for validation, note this only works for image data validationSplit: z.number().min(0).max(1), // batch size of training data diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 5926bbe99..a82b2c69a 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -28,18 +28,6 @@ import { RoundLogs, Trainer } from "./trainer.js"; const debug = createDebug("discojs:training:disco"); -function debugProcessMemory(label: string): void { - if (typeof process === "undefined") return; - - const m = process.memoryUsage(); - debug("%s memory: %O", label, { - rssGB: m.rss / 1024 / 1024 / 1024, - heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, - externalGB: m.external / 1024 / 1024 / 1024, - arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, - }); -} - interface DiscoConfig { scheme: N; logger: Logger; @@ -113,6 +101,7 @@ export class Disco extends EventEmitter<{ readonly #logger: Logger; readonly #task: Task; readonly #preprocessOnce: boolean; + // Forwarded to compatible models to identify this client in debug output. readonly #debugLabel?: string; /** @@ -225,7 +214,6 @@ export class Disco extends EventEmitter<{ } const roundLogs = await roundLogsPromise; - debugProcessMemory(`round ${roundNum} after round logs resolved`); for (const { epochNum, epochLogs } of epochResults) { yield buildSummaryLog(roundNum, epochNum, roundLogs, epochLogs); diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 483675b8f..dc9759289 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -68,10 +68,10 @@ export class Trainer { >; readonly #roundIterations?: number; readonly #validationFrequency?: number; + readonly #validationMode: "before" | "after" | "both"; // Map of weight Index and weight update #weightNormHistory: WeightNormHistory = List(); #previousRoundWeights?: WeightsContainer; - readonly #shouldValidateAfterAggregation: boolean; public get model(): Model { if (this.#model === undefined) @@ -89,8 +89,7 @@ export class Trainer { this.#epochs = task.trainingInformation.epochs; this.#roundIterations = task.trainingInformation.roundIterations; this.#validationFrequency = task.trainingInformation.validationFrequency; - this.#shouldValidateAfterAggregation = - task.trainingInformation.scheme !== "local"; + this.#validationMode = task.trainingInformation.validationMode ?? "before"; if ("privacy" in task.trainingInformation) this.#privacy = task.trainingInformation.privacy; @@ -143,15 +142,19 @@ export class Trainer { try { this.#training = this.#roundIterations === undefined - ? this.#runRounds(dataset, validationDataset) - : this.#runIterationRounds(dataset, validationDataset); + ? this.#runRoundsByEpoch(dataset, validationDataset) + : this.#runRoundsByIteration(dataset, validationDataset); yield* this.#training; } finally { this.#training = undefined; } } - async *#runRounds( + /** + * Runs epoch-based training, aggregating after `roundDuration` complete + * passes over the training dataset until the configured epochs are reached. + */ + async *#runRoundsByEpoch( dataset: Dataset>, validationDataset?: Dataset>, ): AsyncGenerator< @@ -174,13 +177,25 @@ export class Trainer { ? validationDataset : undefined; - yield this.#runRound(dataset, roundValidationDataset); - - await this.#finishRoundCommunication(totalRound); + yield this.#runRoundByEpoch( + dataset, + this.#shouldValidateBeforeAggregation() + ? roundValidationDataset + : undefined, + this.#shouldValidateAfterAggregation() + ? roundValidationDataset + : undefined, + totalRound, + ); } } - async *#runIterationRounds( + /** + * Runs iteration-based training, aggregating after `roundIterations` + * batches while preserving the dataset iterator between rounds. A new + * iterator is created only when the next configured epoch begins. + */ + async *#runRoundsByIteration( dataset: Dataset>, validationDataset?: Dataset>, ): AsyncGenerator< @@ -231,15 +246,19 @@ export class Trainer { ? validationDataset : undefined; - yield this.#runIterationRound( + yield this.#runRoundByIteration( prefixedIterator, this.#roundIterations, - roundValidationDataset, + this.#shouldValidateBeforeAggregation() + ? roundValidationDataset + : undefined, + this.#shouldValidateAfterAggregation() + ? roundValidationDataset + : undefined, + totalRound, (roundDone) => (done = roundDone), ); - await this.#finishRoundCommunication(totalRound); - round++; if (done) break; next = await trainingIterator.next(); @@ -247,9 +266,19 @@ export class Trainer { } } - async *#runRound( + /** + * Trains one epoch-based round by making `roundDuration` complete passes + * over the dataset, then exchanges weights and returns the round metrics. + */ + async *#runRoundByEpoch( dataset: Dataset>, - validationDataset?: Dataset>, + preAggregationValidationDataset: + | Dataset> + | undefined, + postAggregationValidationDataset: + | Dataset> + | undefined, + totalRound: number, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); @@ -257,30 +286,48 @@ export class Trainer { // Before starting the training, get the validation of global model const validation = - validationDataset !== undefined - ? await this.model.evaluate(validationDataset) + preAggregationValidationDataset !== undefined + ? await this.model.evaluate(preAggregationValidationDataset) : undefined; for (let epoch = 0; epoch < this.#roundDuration; epoch++) { const [gen, epochLogs] = async_iterator.split( - this.model.train(dataset, validationDataset), + this.model.train(dataset, preAggregationValidationDataset), ); yield gen; epochsLogs = epochsLogs.push(await epochLogs); } + const participants = this.#client.nbOfParticipants; + const postAggregationValidation = await this.#finishRoundCommunication( + totalRound, + postAggregationValidationDataset, + ); + return { epochs: epochsLogs, - participants: this.#client.nbOfParticipants, + participants, preRoundValidation: validation, + postAggregationValidation, }; } - async *#runIterationRound( + /** + * Trains one iteration-based round by consuming at most `maxBatchCount` + * batches from the supplied iterator without rewinding it, then exchanges + * weights and returns the round metrics. + */ + async *#runRoundByIteration( datasetIterator: AsyncIterator>, maxBatchCount: number, - validationDataset?: Dataset>, + preAggregationValidationDataset: + | Dataset> + | undefined, + postAggregationValidationDataset: + | Dataset> + | undefined, + totalRound: number, setDone?: (done: boolean) => void, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); @@ -288,8 +335,8 @@ export class Trainer { debug("Run iteration-based round"); const validation = - validationDataset !== undefined - ? await this.model.evaluate(validationDataset) + preAggregationValidationDataset !== undefined + ? await this.model.evaluate(preAggregationValidationDataset) : undefined; const model = this.model as unknown as IterationTrainableTextModel; @@ -302,7 +349,7 @@ export class Trainer { Batched >, maxBatchCount, - validationDataset as + preAggregationValidationDataset as | Dataset> | undefined, setDone, @@ -312,13 +359,28 @@ export class Trainer { yield gen; epochsLogs = epochsLogs.push(await epochLogs); + const participants = this.#client.nbOfParticipants; + const postAggregationValidation = await this.#finishRoundCommunication( + totalRound, + postAggregationValidationDataset, + ); + return { epochs: epochsLogs, - participants: this.#client.nbOfParticipants, + participants, preRoundValidation: validation, + postAggregationValidation, }; } + #shouldValidateBeforeAggregation(): boolean { + return this.#validationMode !== "after"; + } + + #shouldValidateAfterAggregation(): boolean { + return this.#validationMode !== "before"; + } + #shouldValidateRound(round: number): boolean { if (this.#validationFrequency === undefined) return true; if (this.#validationFrequency === 0) return false; @@ -366,8 +428,7 @@ export class Trainer { await this.#client.onRoundEndCommunication(roundWeights); this.model.weights = networkWeights; - return this.#shouldValidateAfterAggregation && - validationDataset !== undefined + return validationDataset !== undefined ? await this.model.evaluate(validationDataset) : undefined; } finally { diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index 133650f18..c68e03beb 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -17,20 +17,14 @@ import FederatedMessages = client.federated.messages; const debug = createDebug("server:controllers:federated"); -function debugProcessMemory(label: string): void { - const m = process.memoryUsage(); - debug("%s memory: %O", label, { - rssGB: m.rss / 1024 / 1024 / 1024, - heapUsedGB: m.heapUsed / 1024 / 1024 / 1024, - externalGB: m.external / 1024 / 1024 / 1024, - arrayBuffersGB: m.arrayBuffers / 1024 / 1024 / 1024, - }); -} - export class FederatedController extends TrainingController< D, "federated" > { + /** + * WebSockets of clients whose update was accepted for the current round. + * They receive the resulting global weights when aggregation completes. + */ #pendingUpdateRecipients = new Map(); /** * Aggregators for each hosted task. @@ -52,6 +46,10 @@ export class FederatedController extends TrainingController< this.#latestGlobalWeights = this.initialWeights; } + /** + * Creates an aggregator and registers the handler that caches and broadcasts + * the global weights produced at the end of each aggregation round. + */ #makeAggregator(): aggregators.MeanAggregator { const aggregator = new aggregators.MeanAggregator(undefined, 1, "relative"); @@ -59,14 +57,8 @@ export class FederatedController extends TrainingController< try { const recipients = this.#pendingUpdateRecipients; this.#pendingUpdateRecipients = new Map(); - - debugProcessMemory( - `round ${aggregator.round} before encoding aggregate`, - ); const payload = await serialization.weights.encode(weightUpdate); - debugProcessMemory( - `round ${aggregator.round} after encoding aggregate`, - ); + debug( "round %o aggregate payload byteLength=%d", aggregator.round, @@ -81,31 +73,31 @@ export class FederatedController extends TrainingController< nbOfParticipants: this.connections.size, }; const encodedMsg = msgpack.encode(msg); - debugProcessMemory( - `round ${aggregator.round} after encoding websocket message`, - ); recipients.forEach((recipientWs, recipientId) => { - debug( - "Sending global weights for round %o to client [%s]", - aggregator.round, - recipientId.slice(0, 4), - ); - recipientWs.send(encodedMsg); - debug( - "Aggregated payload sent to client [%s] for round %o", - recipientId.slice(0, 4), - aggregator.round, - ); + try { + debug( + "Sending global weights for round %o to client [%s]", + aggregator.round, + recipientId.slice(0, 4), + ); + recipientWs.send(encodedMsg); + debug( + "Aggregated payload sent to client [%s] for round %o", + recipientId.slice(0, 4), + aggregator.round, + ); + } catch (err) { + debug( + "Failed to send global weights for round %o to client [%s]: %o", + aggregator.round, + recipientId.slice(0, 4), + err, + ); + } }); - debugProcessMemory( - `round ${aggregator.round} after sending aggregate to ${recipients.size} clients`, - ); } finally { weightUpdate.dispose(); - debugProcessMemory( - `round ${aggregator.round} after disposing aggregate tensor`, - ); } }); diff --git a/server/src/routes/training_router.ts b/server/src/routes/training_router.ts index 8c1ff3560..bab92257f 100644 --- a/server/src/routes/training_router.ts +++ b/server/src/routes/training_router.ts @@ -1,5 +1,6 @@ import express from "express"; import type expressWS from "express-ws"; +import createDebug from "debug"; import type { Task, DataType, Network } from "@epfml/discojs"; import { serialization } from "@epfml/discojs"; @@ -10,6 +11,8 @@ import { DecentralizedController, } from "../controllers/index.js"; +const debug = createDebug("server:routes:training"); + /** * The TrainingRouter handles client requests related the federated * and decentralized training. @@ -60,6 +63,9 @@ export class TrainingRouter> { let encodedWeights: serialization.Encoded; try { encodedWeights = await serialization.weights.encode(weights); + } catch (err) { + debug("Failed to encode initial weights for task %s: %o", task.id, err); + return; } finally { model[Symbol.dispose](); } From 757dab10c5afca97c357eb047e1dfdbf9d0c7b11 Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 27 Aug 2026 17:43:26 +0200 Subject: [PATCH 67/89] add avg aggregator change --- discojs/src/aggregator/mean.ts | 21 +++++---------------- discojs/src/aggregator/multiround.ts | 2 +- discojs/src/weights/aggregation.ts | 18 ++++++++++++++++-- 3 files changed, 22 insertions(+), 19 deletions(-) diff --git a/discojs/src/aggregator/mean.ts b/discojs/src/aggregator/mean.ts index 080ff6a03..5b4a59f98 100644 --- a/discojs/src/aggregator/mean.ts +++ b/discojs/src/aggregator/mean.ts @@ -1,10 +1,12 @@ import type { Map } from "immutable"; import { AggregationStep } from "./aggregator.js"; import { MultiRoundAggregator, ThresholdType } from "./multiround.js"; +import { aggregation } from "../index.js"; import type { WeightsContainer, client } from "../index.js"; /** * Mean aggregator whose aggregation step consists in computing the mean of the received weights. + * This aggregator extends MultiRoundAggregator while only performing a single round * */ export class MeanAggregator extends MultiRoundAggregator { @@ -42,22 +44,9 @@ export class MeanAggregator extends MultiRoundAggregator { this.log(AggregationStep.AGGREGATE); const contributions = Array.from(currentContributions.values()); - let summed = contributions[0]?.map((weight) => weight.clone()); - if (summed === undefined) - throw new Error("aggregating without any contribution"); - - try { - for (const contribution of contributions.slice(1)) { - const next = summed.add(contribution); - summed.dispose(); - summed = next; - } - - return summed.map((weight) => weight.div(contributions.length)); - } finally { - summed.dispose(); - contributions.forEach((contribution) => contribution.dispose()); - } + const avg = aggregation.avg(contributions); + contributions.forEach((contribution) => contribution.dispose()); + return avg } override makePayloads( diff --git a/discojs/src/aggregator/multiround.ts b/discojs/src/aggregator/multiround.ts index 4018b55dd..b37aa4853 100644 --- a/discojs/src/aggregator/multiround.ts +++ b/discojs/src/aggregator/multiround.ts @@ -6,7 +6,7 @@ export type ThresholdType = "relative" | "absolute"; const debug = createDebug("discojs:aggregator:multiround"); /** - * Base class for multi-round aggregators. + * Base class for multi-round aggregators, with potentially only one round. * Multi-round aggregators are aggregators that wait for a certain number of contributions before aggregating. * They can be used to implement different aggregation strategies, such as Byzantine robust aggregation or Mean Aggregator. */ diff --git a/discojs/src/weights/aggregation.ts b/discojs/src/weights/aggregation.ts index 868d23b8a..4cf32ee01 100644 --- a/discojs/src/weights/aggregation.ts +++ b/discojs/src/weights/aggregation.ts @@ -67,6 +67,20 @@ export function diff( export function avg( weights: Iterable, ): WeightsContainer { - const ws = List(weights); - return sum(ws).map((w) => w.div(ws.size)); + const ws = parseWeights(weights); + const first = ws.first(); + if (first === undefined) throw new Error("no weights to work with"); + let summed: WeightsContainer = first.map((weight) => weight.clone()); + + try { + for (const weights of ws.rest()) { + const next: WeightsContainer = summed.add(weights); + summed.dispose(); + summed = next; + } + + return summed.map((weight) => weight.div(ws.size)); + } finally { + summed.dispose(); + } } From 3d205c40c24e7dcb4e56b8fc2010cd4662efb697 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 28 Aug 2026 17:26:55 +0200 Subject: [PATCH 68/89] fix: memory leak --- discojs/src/training/trainer.ts | 43 ++++++++++++++++----------------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index dc9759289..d80aade86 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -454,30 +454,9 @@ async function applyOptimalPrivacy( totalRound: number, ): Promise { let ret = current; - - // Clipping radius for BFT - if ("byzantineFaultTolerance" in options) { - // might need to change the variable name - const previousRoundWeights = - previous ?? current.map((w) => tf.zerosLike(w)); - const weightsProgress = current.sub(previousRoundWeights); - const clippedProgress = await privacy.clipNorm( - weightsProgress, - Repeat(options.byzantineFaultTolerance.clippingRadius) - .take(weightsProgress.weights.length) - .toArray(), - ); - try { - ret = previousRoundWeights.add(clippedProgress); - } finally { - weightsProgress.dispose(); - clippedProgress.dispose(); - if (previous === undefined) previousRoundWeights.dispose(); - } - } + const dpOptions = options.differentialPrivacy; // Adding Gaussian noise for DP - const dpOptions = options.differentialPrivacy; if (dpOptions !== undefined) { const dpDefaultRadius = dpOptions.clippingRadius; // options.dpDefaultClippingRadius should be a number @@ -529,5 +508,25 @@ async function applyOptimalPrivacy( if (previous === undefined) previousEpochWeights.dispose(); } } + // Clipping radius for BFT if DP didn't already clip + else if ("byzantineFaultTolerance" in options) { + // might need to change the variable name + const previousRoundWeights = + previous ?? current.map((w) => tf.zerosLike(w)); + const weightsProgress = current.sub(previousRoundWeights); + const clippedProgress = await privacy.clipNorm( + weightsProgress, + Repeat(options.byzantineFaultTolerance.clippingRadius) + .take(weightsProgress.weights.length) + .toArray(), + ); + try { + ret = previousRoundWeights.add(clippedProgress); + } finally { + weightsProgress.dispose(); + clippedProgress.dispose(); + if (previous === undefined) previousRoundWeights.dispose(); + } + } return ret; } From 5fc85dce079da6cb9a8bcf7c52ef2db30724ac12 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 28 Aug 2026 17:27:20 +0200 Subject: [PATCH 69/89] fix: catch potential serialization errors --- server/src/controllers/federated_controller.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index c68e03beb..20375a491 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -55,9 +55,9 @@ export class FederatedController extends TrainingController< aggregator.on("aggregation", async (weightUpdate) => { try { + const payload = await serialization.weights.encode(weightUpdate); const recipients = this.#pendingUpdateRecipients; this.#pendingUpdateRecipients = new Map(); - const payload = await serialization.weights.encode(weightUpdate); debug( "round %o aggregate payload byteLength=%d", @@ -96,6 +96,12 @@ export class FederatedController extends TrainingController< ); } }); + } catch (err) { + debug( + "Failed to serialize or encode weights for round %o: %o", + aggregator.round, + err, + ); } finally { weightUpdate.dispose(); } From 6cc4fb9fc47cbf9aa42da911831ea5877ec460e5 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 28 Aug 2026 17:30:17 +0200 Subject: [PATCH 70/89] fix: outdated command --- cli/package.json | 1 - 1 file changed, 1 deletion(-) diff --git a/cli/package.json b/cli/package.json index b517de895..118d075c5 100644 --- a/cli/package.json +++ b/cli/package.json @@ -12,7 +12,6 @@ "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" }, From 48210656618bdc8d8093bc79e7ce80c54d30648e Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 28 Aug 2026 17:30:26 +0200 Subject: [PATCH 71/89] fix: rm console.log --- decentralized-connection-fix-review.md | 225 +++++++++++++++++++++++++ 1 file changed, 225 insertions(+) create mode 100644 decentralized-connection-fix-review.md diff --git a/decentralized-connection-fix-review.md b/decentralized-connection-fix-review.md new file mode 100644 index 000000000..1c979f762 --- /dev/null +++ b/decentralized-connection-fix-review.md @@ -0,0 +1,225 @@ +# Code Review — `decentralized-connection-fix` + +**Branch:** `decentralized-connection-fix` → `develop` +**Merge base:** `fdabf4ad` +**Scope:** 14 files, ~950 insertions / 27 deletions +**Reviewer:** Claude (Opus 4.8) +**Date:** 2026-08-28 (updated with concrete fixes for M1, M2, L4 and a new design finding M4) + +--- + +## Table of Contents + +1. [Summary](#1-summary) +2. [What the PR does](#2-what-the-pr-does) +3. [Architecture & design notes](#3-architecture--design-notes) +4. [Findings](#4-findings) + - [4.1 High severity](#41-high-severity) + - [4.2 Medium severity](#42-medium-severity) + - [4.3 Low severity / nits](#43-low-severity--nits) +5. [Things done well](#5-things-done-well) +6. [Test coverage](#6-test-coverage) +7. [File-by-file notes](#7-file-by-file-notes) +8. [Recommendation](#8-recommendation) +9. [Appendix: proposed patches](#9-appendix-proposed-patches) + +--- + +## 1. Summary + +This PR makes decentralized learning robust to two real-world conditions that previously broke it: + +1. **Peers joining mid-training** now synchronize the latest model from an existing peer before participating, instead of polluting aggregation with a fresh/untrained model. +2. **Peer-connection failures** are now handled with a server-coordinated retry loop, a connection-readiness barrier (so fast peers don't start weight sharing before slow peers connect), and graceful exclusion of peers that never connect. + +The change is well-scoped, thoroughly documented (a new [`README.md`](discojs/src/client/decentralized/README.md) explains the whole protocol), and backed by substantial new e2e tests. The protocol design is sound. My concerns are mostly about a few edge-case correctness bugs and some listener/connection hygiene rather than the overall approach. + +**Overall: approve with changes requested.** The High/Medium items below should be addressed before merge; the rest can be follow-ups. + +--- + +## 2. What the PR does + +New message types added across `messages.ts` (both shared and decentralized): + +- `ConnectionsReady` (peer → server): "I've established all my round peer connections." +- `StartWeightSharing` (server → peers): readiness barrier release. +- `RetryPeerConnections` / `ConnectionFail` (server → peers): retry/exclusion control. +- `ModelSyncRequest` (peer → server), `SignalModelProvider` / `SignalNewPeer` (server → peers), `SharedModel` (peer → peer): mid-training model sync handshake. + +Supporting changes: + +- `Disco` injects a `ModelWeightAccess` interface into the client so the client can read/write trainer weights for syncing (`discojs/src/training/disco.ts`). See **M4** for a design concern with this seam. +- `Trainer` now calls `client.finishRound()` after applying aggregated weights (`discojs/src/training/trainer.ts`). +- New `maxConnectionRetry` task field with a Zod schema + task-creation form field, defaulting to 3. +- UI toasts for the new failure modes in `TrainerDashboard.vue`. + +--- + +## 3. Architecture & design notes + +The protocol is layered cleanly on top of the existing server-coordinated signaling: + +- The **readiness barrier** (`ConnectionsReady` → `StartWeightSharing`) is a good fix for the documented race where a fast peer begins weight sharing while a slow peer is still negotiating WebRTC. +- **Model sync reuses the existing WebRTC signaling path** (`SignalForPeer`), with the server brokering a provider/newcomer pairing. This is the right call — no new transport. +- The **provider waits for the in-flight aggregation round to finish** (`roundFinishedPromise`) before sending weights, so the newcomer gets a consistent aggregated model. Because all peers converge to the same aggregated weights each round, sending from a single provider is sufficient. Good. + +The main design risk is **shared mutable state reuse**: the model-sync connections are created on the *same* `PeerPool` and the *same* server-message stream used for normal rounds. This is where most of the findings below originate. + +--- + +## 4. Findings + +### 4.1 High severity + +#### H1 — Provider re-selection on disconnect can pick a syncing (stale-model) node + +`server/src/controllers/decentralized_controller.ts:164` + +```ts +if (this.#providerNode === peerId) { + this.#providerNode = this.connections.keySeq().first() +} +``` + +When the current provider disconnects, the replacement is just the first connection in the map. That candidate may itself be in `#syncingNodes` — i.e. a peer that joined mid-training and has **not yet received the latest model**. If a `ModelSyncRequest` then routes to it, `handleSignalNewPeer` will read its un-synced weights via `getModelWeight()` and ship a stale/wrong model to the newcomer, silently corrupting its starting point. + +**Fix:** pick the replacement from `connections.keySeq().toSet().subtract(this.#syncingNodes)` (and handle the empty case by clearing `#providerNode`). + +--- + +### 4.2 Medium severity + +#### M1 — Leaked `once` listeners from `Promise.race(waitMessage(...))` + +`discojs/src/client/decentralized/decentralized_client.ts:238` + +```ts +const msg = await Promise.race([ + waitMessage(this.server, type.StartWeightSharing), + waitMessage(this.server, type.RetryPeerConnections), + waitMessage(this.server, type.ConnectionFail), +]) +``` + +`waitMessage` registers a `once` listener that is only removed when *that specific event type* is emitted (see `EventEmitter.emit` in `discojs/src/utils/event_emitter.ts` — it filters out `once` listeners only for the emitted event). When the race settles on, say, `StartWeightSharing`, the `RetryPeerConnections` and `ConnectionFail` listeners stay registered. Every round and every `continue` of the retry loop adds three more, so listeners accumulate for the lifetime of the connection. + +Correctness is mostly preserved (the orphaned resolves are ignored by the already-settled `Promise.race`), but it's a genuine listener leak that grows unbounded over a long training run, and when a `RetryPeerConnections`/`ConnectionFail` finally fires it wakes every stale listener. + +**Same bug exists in `waitMessageWithTimeout`** (`event_connection.ts:30`): it does `Promise.race([waitMessage(...), timeout(...)])`, so when the timeout wins the `waitMessage` listener leaks. That path backs the 30 s `SignalModelProvider` and `SharedModel` sync waits, so fix it in the same change. + +**Fix (concrete):** the root cause is that `EventEmitter` has no way to remove a listener. Add `off`, then introduce a single cancel-clean primitive and route the existing helpers + the race through it: + +1. `EventEmitter.off(event, listener)` in `event_emitter.ts` (reference-equality removal — safe because each waiter uses fresh closures). +2. Declare `off` on the `EventConnection` interface (both `PeerConnection` and `WebSocketServer` inherit it from `EventEmitter`). +3. Add `waitForFirstMessage(connection, types[], {timeoutMs?, errorMsg?})` that registers via `on`, and in a `finally` removes **every** listener and clears the timer. Re-express `waitMessage` / `waitMessageWithTimeout` in terms of it. +4. Replace the race with `await waitForFirstMessage(this.server, [StartWeightSharing, RetryPeerConnections, ConnectionFail])` — the `msg.type` discrimination downstream is unchanged. + +This closes both the round-control leak and the latent `waitMessageWithTimeout` leak in one place. Full patch in the [appendix](#9-appendix-proposed-patches). + +#### M2 — Sync connections leak into the shared `PeerPool` + +`discojs/src/client/decentralized/decentralized_client.ts:52` and `:182` + +`PeerPool.getPeers` merges new connections into the shared `this.peers` (`discojs/src/client/decentralized/peer_pool.ts:63`) and never removes them. The model-sync path calls `getPeers(Set([newNode]))` / `getPeers(Set([providerNode]))` on the same `#pool` used for normal rounds. After sync, that provider/newcomer connection lingers in the pool indefinitely (it's filtered out of the round's return value but stays resident). It isn't added to the aggregator node set, so aggregation stays correct, but it's an un-cleaned WebRTC connection per mid-training join. + +**Fix:** either tear down the sync connection once `sendModel`/`receiveModel` completes, or use a dedicated short-lived pool for syncing. + +#### M3 — Retry-loop reset doesn't refresh `aggregationResult` + +`discojs/src/client/decentralized/decentralized_client.ts:205` vs `:259` + +`aggregationResult` is captured once in `onRoundBeginCommunication` (`this.aggregator.getPromiseForAggregation()`). On a `RetryPeerConnections`, the loop resets the pool, `#connections`, and `setAggregatorNodes(Set(this.ownId))`, then loops back to `establishPeerConnections`. The `aggregationResult` promise from round-begin is **not** re-created after the aggregator node set is reset, yet `exchangeWeightUpdates` later awaits it. Worth confirming with a forced-retry test that the aggregator promise created before the node-set reset still resolves correctly after the reset; if not, re-capture it after a successful (post-retry) `establishPeerConnections`. + +--- + +### 4.3 Low severity / nits + +#### L1 — `RetryPeerConnections` interface field is never populated + +The interface declares `aggregationRound: number` (`messages.ts:54`), but the server sends `{ type: MessageTypes.RetryPeerConnections }` without it (`decentralized_controller.ts:323`, `:358`), and the client never reads it. Either populate it or drop the field. (TypeScript doesn't catch this because the server uses an untyped object literal.) + +#### L2 — Inconsistent message validation + +In `isMessageFromServer` (`messages.ts`), `SignalNewPeer` validates its `newNode` payload, but `SignalModelProvider` (which also carries a `providerNode: NodeID`) falls into the catch-all `return true` group with no field check (`messages.ts:146`). Validate `providerNode` for symmetry/robustness. + +#### L3 — Misleading indentation in `handleSignalNewPeer` + +`decentralized_client.ts:49-60`: the body after the brace-less `if (this.#pool === undefined) throw ...` is indented as if it were inside the `if`. It executes unconditionally (which is correct), but the indentation reads as a bug. Add braces or de-indent. + +#### L4 — Duplicated retry-broadcast block + +`decentralized_controller.ts` `handleTimeout` repeats the same "map → encode → lookup conn → send" broadcast almost verbatim in both the exclusion branch and the plain-retry branch (and again in `signalWeightSharing` / `sendPeersForRoundIfNeeded`). Extract a `broadcast(nodeIds, message)` helper to cut ~60 lines and the risk of the branches drifting apart. + +#### L5 — `modelSynced` event can emit `undefined` + +`decentralized_client.ts:196`: `this.emit("modelSynced", this.modelWeightAccess?.getModelWeight())` emits `undefined` when `modelWeightAccess` is unset, and the event type is `WeightsContainer | undefined`. Every consumer (including the e2e test) has to guard against `undefined`. Consider only emitting when weights exist, and typing the event as `WeightsContainer`. + +#### L6 — Off-by-one between README and retry threshold + +README step says "if retries still **below** `maxConnectionRetry`, send Retry," but the code excludes when `#connectionRetry >= maxConnectionRetry` after incrementing first (`decentralized_controller.ts:299`). With `maxConnectionRetry = 3` you get retries on attempts 1 and 2, exclusion on attempt 3. Functionally fine; tighten the wording so the count matches. + +#### L7 — Slow-but-healthy peers can be excluded as "failed" + +`handleTimeout` treats any node that hasn't sent `ConnectionsReady` within the timeout as failed and sends it `ConnectionFail` → disconnect. A genuinely slow peer (large model, slow NAT traversal) gets evicted. The 60s timeout is generous, but worth a comment noting the trade-off, and possibly making the timeout configurable alongside `maxConnectionRetry`. + +#### L8 — Hardcoded 30s sync timeouts + +`receiveModel` / `onRoundBeginCommunication` use literal `30_000` (`decentralized_client.ts`). For large models (e.g. the wikitext/GPT-2 tasks) 30s for a full weight transfer over WebRTC may be tight. Consider deriving from task config or at least a named constant. + +--- + +## 5. Things done well + +- **Excellent protocol documentation** — `README.md` walks through connection, readiness barrier, retry/exclusion, and mid-training sync step-by-step. This is the kind of doc that makes the feature maintainable. +- **The `aggregationStrategy` swap in `TaskCreationForm.vue` is a genuine fix, not a regression.** The form previously offered `secure` for *federated* and withheld it from *decentralized*; this PR aligns it with both `training_information.ts` and `aggregator/get.ts:75` ("secure aggregation is currently supported for decentralized only"). ✅ +- **Exhaustiveness preserved** — the server's `const _: never = msg` switch guard still holds after adding `ConnectionsReady` and `ModelSyncRequest` handlers. +- **`ModelWeightAccess` abstraction** keeps the client decoupled from `Trainer` internals — clean seam. +- **Base-class `finishRound()` no-op** means the federated client is unaffected by the new trainer hook. Good defensive default. +- **User-facing error toasts** for the two new failure modes, with actionable "please rejoin" messaging. + +--- + +## 6. Test coverage + +`server/tests/e2e/decentralized.spec.ts` adds strong coverage: + +- 10-participant consensus tests for **mean**, **byzantine**, and **secure** aggregators — exercises the readiness barrier at non-trivial scale. +- A dedicated **mid-training join / model-sync** test that starts User3 after User1/User2 have completed a round and asserts User3's synced weights exactly equal a provider's weights. + +Gaps worth adding: + +- **No test forces a connection failure / retry / exclusion path** (`RetryPeerConnections`, `ConnectionFail`, `maxConnectionRetry`). Given the complexity of `handleTimeout`, this is the riskiest untested code in the PR (and relates to M3). A test that drops/stalls one peer's connection would be valuable. +- The mid-training test's `for (attempt … 5)` + `setTimeout` polling is inherently timing-dependent and may flake under CI load; consider driving it deterministically via the `modelSynced` event alone. + +--- + +## 7. File-by-file notes + +| File | Notes | +|------|-------| +| `discojs/src/client/decentralized/decentralized_client.ts` | Core of the change. See H1-context, M1, M2, M3, L3, L5, L8. | +| `server/src/controllers/decentralized_controller.ts` | Server orchestration + retry. See H1, L1, L4, L6, L7. | +| `discojs/src/client/decentralized/messages.ts` | New message types; see L1, L2. | +| `discojs/src/client/messages.ts` | New `type` enum entries — clean, well-commented. | +| `discojs/src/client/client.ts` | Adds `setModelWeightAccess`, `finishRound()` no-op, `modelSynced` event. Fine. Minor: mixed tab/space indentation around the new method. | +| `discojs/src/training/disco.ts` | `ModelWeightAccess` wiring. `getModelWeight` clones tensors every call — acceptable; just be aware of tfjs tensor disposal over long runs. | +| `discojs/src/training/trainer.ts` | One-line `finishRound()` hook. Correct placement (after weights applied). | +| `discojs/src/task/training_information.ts` | `maxConnectionRetry` Zod field with sane default. Good. | +| `discojs/src/default_tasks/{cifar10,mnist}.ts` | Add `maxConnectionRetry: 3`. Note other decentralized default tasks aren't updated but rely on the schema default — verify that's intended. | +| `webapp/src/components/task_creation_form/TaskCreationForm.vue` | New retry field + the correct strategy swap (see §5). Stray trailing whitespace/blank lines after the new `FormLabel`. | +| `webapp/src/components/training/TrainerDashboard.vue` | New toasts + `trainingCompleted` guard so success isn't shown after a handled failure. Typo: "compeleted" in comment. | + +--- + +## 8. Recommendation + +**Approve with changes requested.** + +Must-fix before merge: +- **H1** (stale-provider selection on disconnect) — correctness bug that can silently corrupt a newcomer's model. + +Strongly recommended before merge: +- **M1** (listener leak), **M2** (pool connection leak), **M3** (verify `aggregationResult` after retry — at minimum add a retry-path e2e test). + +Everything in §4.3 is fine as follow-up. The design is solid and the documentation/tests are above the bar for this codebase — nice work. From 99abb0ae13c57bde87d62b92c2b83ff6009c7b91 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Fri, 28 Aug 2026 17:33:16 +0200 Subject: [PATCH 72/89] fix: only log when metric is defined --- discojs/src/training/disco.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index a82b2c69a..a2ffe381a 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -309,8 +309,10 @@ export class Disco extends EventEmitter<{ this.#logger.success( [ `Round: ${roundNum}`, - `Post-aggregation loss: ${roundLogs.postAggregationValidation?.loss}`, - `Post-aggregation accuracy: ${roundLogs.postAggregationValidation?.accuracy}`, + roundLogs.postAggregationValidation !== undefined ? + `Post-aggregation loss: ${roundLogs.postAggregationValidation.loss}` : "", + roundLogs.postAggregationValidation ? + `Post-aggregation accuracy: ${roundLogs.postAggregationValidation.accuracy}` : "", ].join("\n"), ); From 3b917a614d30feebd271b35ced716a72a12d0c4c Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:18:46 +0200 Subject: [PATCH 73/89] fix: rm console.log --- cli/src/measure_memorization_gpt2.ts | 47 ++++++++++++++-------------- 1 file changed, 24 insertions(+), 23 deletions(-) diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts index 165beb589..68cf5bcfc 100644 --- a/cli/src/measure_memorization_gpt2.ts +++ b/cli/src/measure_memorization_gpt2.ts @@ -463,29 +463,30 @@ async function main() { prompt.length, prompt.length + args.suffixLength, ); - - console.log("================================"); - console.log("PROMPT LENGTH:", promptLength); - - console.log("\nPROMPT IDS:"); - console.log(prompt.slice(0, 30)); - - console.log("\nGENERATED IDS:"); - console.log(generatedSuffix.slice(0, 30)); - - console.log("\nREFERENCE IDS:"); - console.log(reference.slice(0, 30)); - - console.log("\nPROMPT TEXT:"); - console.log(JSON.stringify(tokenizer.decode(prompt))); - - console.log("\nGENERATED TEXT:"); - console.log(JSON.stringify(tokenizer.decode(generatedSuffix))); - - console.log("\nREFERENCE TEXT:"); - console.log(JSON.stringify(tokenizer.decode(reference))); - - console.log("================================"); + if (shouldLogRecord) { + console.log("================================"); + console.log("PROMPT LENGTH:", promptLength); + + console.log("\nPROMPT IDS:"); + console.log(prompt.slice(0, 30)); + + console.log("\nGENERATED IDS:"); + console.log(generatedSuffix.slice(0, 30)); + + console.log("\nREFERENCE IDS:"); + console.log(reference.slice(0, 30)); + + console.log("\nPROMPT TEXT:"); + console.log(JSON.stringify(tokenizer.decode(prompt))); + + console.log("\nGENERATED TEXT:"); + console.log(JSON.stringify(tokenizer.decode(generatedSuffix))); + + console.log("\nREFERENCE TEXT:"); + console.log(JSON.stringify(tokenizer.decode(reference))); + + console.log("================================"); + } const exactMatch = generatedSuffix.length === reference.length && From 6f644a90346ffb1b47466b1bcc7cda2c310f31de Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:19:58 +0200 Subject: [PATCH 74/89] fix: rm unreachable if --- discojs/src/training/disco.ts | 5 ----- 1 file changed, 5 deletions(-) diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index a2ffe381a..6f11dd747 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -256,11 +256,6 @@ export class Disco extends EventEmitter<{ this.#setModelDebugLabel(this.trainer.model); this.#setModelTrainingOptions(this.trainer.model); debug("Initial model fetched successfully"); - if (this.trainer.model === null) { - debug( - `No pre-trained model provided for client, initializing randomly...`, - ); - } for await (const [roundNum, round] of enumerate( this.trainer.train(trainingDataset, validationDataset_), From c3423ac6cf4bfaec397470f5b30855718bcba16d Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:20:18 +0200 Subject: [PATCH 75/89] fix: add validationMode arg to cli --- cli/src/args.ts | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/cli/src/args.ts b/cli/src/args.ts index 61d326ac6..002640c81 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -11,6 +11,16 @@ 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; testID: string; @@ -32,6 +42,12 @@ export interface BenchmarkArguments { * `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 @@ -127,6 +143,13 @@ const unsafeArgs = parse( "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", @@ -308,6 +331,7 @@ export const args: BenchmarkArguments = { task.trainingInformation.roundIterations = unsafeArgs.roundIterations; task.trainingInformation.validationFrequency = unsafeArgs.validationFrequency; + task.trainingInformation.validationMode = unsafeArgs.validationMode; if (unsafeArgs.goldfishLoss) { if ( From c564c60ee1c257901e054c4134463780825982db Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:20:37 +0200 Subject: [PATCH 76/89] fix: propagate encoding error in catch --- server/src/routes/training_router.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/routes/training_router.ts b/server/src/routes/training_router.ts index bab92257f..4e61315ce 100644 --- a/server/src/routes/training_router.ts +++ b/server/src/routes/training_router.ts @@ -65,7 +65,7 @@ export class TrainingRouter> { encodedWeights = await serialization.weights.encode(weights); } catch (err) { debug("Failed to encode initial weights for task %s: %o", task.id, err); - return; + throw err; } finally { model[Symbol.dispose](); } From 50f5174a4301e36536c0e4b3ccc77a5183615c88 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:27:10 +0200 Subject: [PATCH 77/89] fix: change console.log to debug --- discojs/src/training/disco.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 6f11dd747..b62546c0a 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -347,7 +347,7 @@ export class Disco extends EventEmitter<{ ); if (this.#task.trainingInformation.goldfishLoss?.enabled === true) { const { k, h, padTokenId } = this.#task.trainingInformation.goldfishLoss; - console.log( + debug( `Using Goldfish loss with k=${k}, h=${h}` + (padTokenId === undefined ? "" : `, padTokenId=${padTokenId}`), ); @@ -356,7 +356,7 @@ export class Disco extends EventEmitter<{ configurableModel.setLearningRate?.( this.#task.trainingInformation.learningRate, ); - console.log( + debug( `Using GPT learning rate ${this.#task.trainingInformation.learningRate}`, ); } From 4fb39d8c92d3912987d011401a95ded031fc522f Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:28:44 +0200 Subject: [PATCH 78/89] fix: remove debug statements --- discojs/src/client/client.ts | 7 ------- 1 file changed, 7 deletions(-) diff --git a/discojs/src/client/client.ts b/discojs/src/client/client.ts index 9577397a2..3f881ea9c 100644 --- a/discojs/src/client/client.ts +++ b/discojs/src/client/client.ts @@ -198,15 +198,8 @@ export abstract class Client extends EventEmitter<{ } url.pathname += `tasks/${this.task.id}/model.json`; - debug( - "fetching latest model from server at %0 for task %1...", - url.href, - this.task.id, - ); - const response = await fetch(url); if (!response.ok) throw new Error(`fetch: HTTP status ${response.status}`); - else debug("response ok, decoding model..."); const encoded = new Uint8Array(await response.arrayBuffer()); return await serialization.model.decode(encoded); From b2ee266d10b859aaf60fe35077b8059602c62c53 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:37:15 +0200 Subject: [PATCH 79/89] fix: throw error if validation path is set for non-text task --- cli/src/data.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cli/src/data.ts b/cli/src/data.ts index 1bd68e0d7..204567712 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -154,6 +154,9 @@ export async function getTaskData( isValidation?: boolean, validationDatasetPath?: string, ): Promise> { + if (validationDatasetPath && taskID !== "privacyrun" && taskID !== "centralized-gpt2-finetune") + throw new Error("validationDatasetPath is currently only supported for text tasks") + switch (taskID) { case "simple_face": // remove return (await loadSimpleFaceData(userIdx, totalClient)) as Dataset< From da5896343641ffcae5e3b80254487af68cbdffbb Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:37:24 +0200 Subject: [PATCH 80/89] fix: outdated doc --- discojs/src/default_tasks/centralized_gpt2_finetune.ts | 2 +- discojs/src/default_tasks/privacyrun.ts | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/discojs/src/default_tasks/centralized_gpt2_finetune.ts b/discojs/src/default_tasks/centralized_gpt2_finetune.ts index ea6a8df5a..4c0e60508 100644 --- a/discojs/src/default_tasks/centralized_gpt2_finetune.ts +++ b/discojs/src/default_tasks/centralized_gpt2_finetune.ts @@ -18,7 +18,7 @@ export const centralizedGPT2FineTune: TaskProvider<"text", "local"> = { "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", "The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer.", "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", - "Context length is kept at 1024 to match the pre-trained model, with batch size at 1.", + "Context length is kept at 512 to match the pre-trained model, with batch size at 8.", ].join(" "), dataFormatInformation: "You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.", diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/privacyrun.ts index 3bd0cc984..da9ae937f 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/privacyrun.ts @@ -18,7 +18,7 @@ export const privacyrun: TaskProvider<"text", "federated"> = { "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", "The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer.", "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", - "Context length is kept at 1024 to match the pre-trained model, with batch size at 1.", + "Context length is kept at 512 to match the pre-trained model, with batch size at 8.", ].join(" "), dataFormatInformation: "You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.", From ccdbca0f28df854065e26cff88d5f01cc6c109c7 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 15:43:40 +0200 Subject: [PATCH 81/89] test: aggregator tensor disposal --- discojs/src/aggregator/mean.spec.ts | 24 ++++++++++++ discojs/src/weights/aggregation.spec.ts | 51 +++++++++++++++++++++++++ 2 files changed, 75 insertions(+) diff --git a/discojs/src/aggregator/mean.spec.ts b/discojs/src/aggregator/mean.spec.ts index c9eab512f..6602e14b9 100644 --- a/discojs/src/aggregator/mean.spec.ts +++ b/discojs/src/aggregator/mean.spec.ts @@ -1,3 +1,4 @@ +import * as tf from "@tensorflow/tfjs"; import { Set } from "immutable"; import { describe, expect, it } from "vitest"; import { WeightsContainer } from "../index.js"; @@ -65,6 +66,29 @@ describe("mean aggregator", () => { expect(await WSIntoArrays(await results)).to.deep.equal([[1], [2]]); }); + it("aggregation leaves no dangling tensors", async () => { + const baseline = tf.memory().numTensors; + + const aggregator = new MeanAggregator(0, 2, "absolute"); + const [id1, id2] = ["client 1", "client 2"]; + aggregator.setNodes(Set.of(id1, id2)); + + const contribution1 = WeightsContainer.of([0], [1]); + const contribution2 = WeightsContainer.of([2], [3]); + + const result = aggregator.getPromiseForAggregation(); + aggregator.add(id1, contribution1, 0); + aggregator.add(id2, contribution2, 0); + expect(await WSIntoArrays(await result)).to.deep.equal([[1], [2]]); + + // the aggregator clones contributions on add and disposes the clones when + // aggregating; only the caller-owned inputs and the result should remain + contribution1.dispose(); + contribution2.dispose(); + (await result).dispose(); + expect(tf.memory().numTensors).to.equal(baseline); + }); + it("waits for 100% of the contributions by default", async () => { const aggregator = new MeanAggregator(); const [id1, id2] = ["client 1", "client 2"]; diff --git a/discojs/src/weights/aggregation.spec.ts b/discojs/src/weights/aggregation.spec.ts index b9757be71..94be3ec7a 100644 --- a/discojs/src/weights/aggregation.spec.ts +++ b/discojs/src/weights/aggregation.spec.ts @@ -1,3 +1,4 @@ +import * as tf from "@tensorflow/tfjs"; import { assert, describe, it } from "vitest"; import { WeightsContainer, aggregation } from "./index.js"; @@ -13,6 +14,56 @@ describe("weights aggregation", () => { assert.isTrue(actual.equals(expected)); }); + it("avg does not leak intermediate tensors", () => { + const baseline = tf.memory().numTensors; + + const inputs = [ + WeightsContainer.of([1, 2, 3, -1], [-5, 6]), + WeightsContainer.of([2, 3, 7, 1], [-10, 5]), + WeightsContainer.of([3, 1, 5, 3], [-15, 19]), + ]; + const result = aggregation.avg(inputs); + + inputs.forEach((input) => input.dispose()); + result.dispose(); + assert.strictEqual(tf.memory().numTensors, baseline); + }); + + it("avg of a single container does not leak intermediate tensors", async () => { + const baseline = tf.memory().numTensors; + + const input = WeightsContainer.of([1, 2], [3]); + const result = aggregation.avg([input]); + + // read values without allocating comparison tensors + assert.deepStrictEqual( + await Promise.all(result.weights.map(async (w) => [...(await w.data())])), + [[1, 2], [3]], + ); + input.dispose(); + result.dispose(); + assert.strictEqual(tf.memory().numTensors, baseline); + }); + + it("avg does not dispose nor alias its inputs", async () => { + const inputs = [ + WeightsContainer.of([1, 2], [3]), + WeightsContainer.of([3, 4], [5]), + ]; + const result = aggregation.avg(inputs); + + for (const input of inputs) + for (const weight of input.weights) assert.isFalse(weight.isDisposed); + + // the result must own fresh tensors, not views of the inputs + inputs.forEach((input) => input.dispose()); + assert.deepStrictEqual( + await Promise.all(result.weights.map(async (w) => [...(await w.data())])), + [[2, 3], [4]], + ); + result.dispose(); + }); + it("sum of weights with two operands", () => { const actual = aggregation.sum([ [[3, -4], [9]], From 019fb9b59b39755ca1f3976712b52b3d92d8e006 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:18:58 +0200 Subject: [PATCH 82/89] test: tensor disposal --- discojs/src/privacy.spec.ts | 45 +++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/discojs/src/privacy.spec.ts b/discojs/src/privacy.spec.ts index 451a7ea78..8a3203e1c 100644 --- a/discojs/src/privacy.spec.ts +++ b/discojs/src/privacy.spec.ts @@ -23,6 +23,17 @@ describe("frobeniusNorm", () => { const n = await frobeniusNorm(t); expect(n).toBeCloseTo(5, 1e-12); }); + + it("does not leak intermediate tensors nor dispose its input", async () => { + const baseline = tf.memory().numTensors; + + const t = tf.tensor([3, 4]); + await frobeniusNorm(t); + + expect(t.isDisposed).toBe(false); + t.dispose(); + expect(tf.memory().numTensors).toBe(baseline); + }); }); describe("clipNorm", () => { @@ -46,6 +57,22 @@ describe("clipNorm", () => { [0, 3], ]); }); + + it("does not leak intermediate tensors nor dispose its input", async () => { + const baseline = tf.memory().numTensors; + + // one layer above the radius (clipped), one within it (kept as-is) + const input = WeightsContainer.of([2], [0, 6]); + const result = await clipNorm(input, [1, 10]); + + for (const weight of input.weights) expect(weight.isDisposed).toBe(false); + + // the result must own fresh tensors, not views of the input + input.dispose(); + expect(await WSIntoArrays(result)).toEqual([[1], [0, 6]]); + result.dispose(); + expect(tf.memory().numTensors).toBe(baseline); + }); }); describe("addOptimalNoise", () => { @@ -69,6 +96,24 @@ describe("addOptimalNoise", () => { expect(Number.isFinite(resultArrays[1][0])).toBe(true); expect(Number.isFinite(resultArrays[1][1])).toBe(true); }); + + it("does not leak intermediate tensors nor dispose its input", async () => { + const baseline = tf.memory().numTensors; + + const input = WeightsContainer.of([3, 4], [0, 6]); + const result = await addOptimalNoise(input, 1, 1e-5, [5, 3]); + + for (const weight of input.weights) expect(weight.isDisposed).toBe(false); + + // the internally clipped weights must be disposed and the result must own + // fresh tensors, not views of the input + input.dispose(); + expect((await WSIntoArrays(result)).flat().every(Number.isFinite)).toBe( + true, + ); + result.dispose(); + expect(tf.memory().numTensors).toBe(baseline); + }); }); describe("getClippingRadius", () => { From 21f640e4d347912206f4df108af949d0c83d1b69 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:19:12 +0200 Subject: [PATCH 83/89] fix: rm debug statements --- discojs/src/models/gpt/index.ts | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/discojs/src/models/gpt/index.ts b/discojs/src/models/gpt/index.ts index e75dda4e1..95600d223 100644 --- a/discojs/src/models/gpt/index.ts +++ b/discojs/src/models/gpt/index.ts @@ -293,17 +293,8 @@ export class GPT extends Model<"text"> { } static deserialize(data: GPTSerialization): Model<"text"> { - debug("GPT model deserialization started"); - - const config = data.config; - - const model = new GPT(config); - - debug("GPT model config initialized: %O", config); - + const model = new GPT(data.config); model.weights = data.weights; - - debug("GPT model weights initialized"); return model; } From 036dbdbe68124770e62e2febb26e3669a996d9b7 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:19:36 +0200 Subject: [PATCH 84/89] doc: roundIterations takes precedence over roundDuration --- discojs/src/task/training_information.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index 8251d9872..221919cd8 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -67,6 +67,7 @@ export namespace TrainingInformation { // e.g.if 3 then weights are shared every 3 epochs (in the distributed setting). roundDuration: z.number().positive().int(), // for GPT text tasks, number of training batches between each weight sharing round. + // roundDuration is ignored if roundIterations is set roundIterations: z.number().positive().int().optional(), // run validation every N aggregation rounds. If 0, validation metrics are skipped. validationFrequency: z.number().nonnegative().int().optional(), From 56633db60660c20696d42eb5295085368570b187 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:20:06 +0200 Subject: [PATCH 85/89] fix: throw early if model doesn't support training by iteration --- discojs/src/training/trainer.ts | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index d80aade86..09d8bd674 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -330,19 +330,18 @@ export class Trainer { totalRound: number, setDone?: (done: boolean) => void, ): AsyncGenerator, RoundLogs> { - let epochsLogs = List(); - + const model = this.model as unknown as IterationTrainableTextModel; + if (typeof model.trainNextBatches !== "function") + throw new Error("model does not support iteration-based training"); + debug("Run iteration-based round"); + let iterationLogs = List(); const validation = preAggregationValidationDataset !== undefined ? await this.model.evaluate(preAggregationValidationDataset) : undefined; - const model = this.model as unknown as IterationTrainableTextModel; - if (typeof model.trainNextBatches !== "function") - throw new Error("model does not support iteration-based training"); - const [gen, epochLogs] = async_iterator.split( model.trainNextBatches( datasetIterator as AsyncIterator< @@ -357,7 +356,7 @@ export class Trainer { ); yield gen; - epochsLogs = epochsLogs.push(await epochLogs); + iterationLogs = iterationLogs.push(await epochLogs); const participants = this.#client.nbOfParticipants; const postAggregationValidation = await this.#finishRoundCommunication( @@ -366,7 +365,7 @@ export class Trainer { ); return { - epochs: epochsLogs, + epochs: iterationLogs, participants, preRoundValidation: validation, postAggregationValidation, From d0db51b0c59bba7150ca5e7ea0e58f94647e6320 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:29:26 +0200 Subject: [PATCH 86/89] refactor: rename privacyrun to goldfish --- cli/src/args.ts | 2 +- cli/src/data.ts | 6 +++--- .../default_tasks/{privacyrun.ts => goldfish.ts} | 14 +++++++++----- discojs/src/default_tasks/index.ts | 2 +- 4 files changed, 14 insertions(+), 10 deletions(-) rename discojs/src/default_tasks/{privacyrun.ts => goldfish.ts} (92%) diff --git a/cli/src/args.ts b/cli/src/args.ts index 002640c81..9afcfc391 100644 --- a/cli/src/args.ts +++ b/cli/src/args.ts @@ -300,7 +300,7 @@ const supportedTasks = Map( defaultTasks.titanic, defaultTasks.tinderDog, defaultTasks.mnist, - defaultTasks.privacyrun, + defaultTasks.goldfish, defaultTasks.centralizedGPT2FineTune, ).map( async (t) => diff --git a/cli/src/data.ts b/cli/src/data.ts index 204567712..2ac6f79d0 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -154,9 +154,9 @@ export async function getTaskData( isValidation?: boolean, validationDatasetPath?: string, ): Promise> { - if (validationDatasetPath && taskID !== "privacyrun" && taskID !== "centralized-gpt2-finetune") + if (validationDatasetPath && taskID !== "goldfish" && taskID !== "centralized-gpt2-finetune") throw new Error("validationDatasetPath is currently only supported for text tasks") - + switch (taskID) { case "simple_face": // remove return (await loadSimpleFaceData(userIdx, totalClient)) as Dataset< @@ -183,7 +183,7 @@ export async function getTaskData( case "mnist_federated": case "mnist": return loadData("mnist", userIdx) as Dataset; - case "privacyrun": + case "goldfish": case "centralized-gpt2-finetune": { const filePath = isValidation && validationDatasetPath diff --git a/discojs/src/default_tasks/privacyrun.ts b/discojs/src/default_tasks/goldfish.ts similarity index 92% rename from discojs/src/default_tasks/privacyrun.ts rename to discojs/src/default_tasks/goldfish.ts index da9ae937f..f4e31f55b 100644 --- a/discojs/src/default_tasks/privacyrun.ts +++ b/discojs/src/default_tasks/goldfish.ts @@ -1,16 +1,16 @@ import type { TaskProvider } from "../index.js"; import { Tokenizer, models, serialization } from "../index.js"; -export const privacyrun: TaskProvider<"text", "federated"> = { +export const goldfish: TaskProvider<"text", "federated"> = { async getTask() { return { - id: "privacyrun", + id: "goldfish", dataType: "text", displayInformation: { - title: "GPT Privacy-Preserving Fine-tuning", + title: "Privacy-Preserving Fine-tuning of GPT-2", summary: { preview: - "Fine-tune a pre-trained GPT model collaboratively and privately.", + "Fine-tune a pre-trained GPT model collaboratively and privately with the Goldfish loss.", overview: "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning.", }, @@ -35,9 +35,13 @@ export const privacyrun: TaskProvider<"text", "federated"> = { // Last context segment may be shorter than context length, so it will be dropped (TODO: implement padding to avoid this) batchSize: 8, tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), - // contextLength: 1024, contextLength: 512, tensorBackend: "gpt", + goldfishLoss: { + enabled: true, + k: 4, + h: 13, + } }, }; }, diff --git a/discojs/src/default_tasks/index.ts b/discojs/src/default_tasks/index.ts index 1ca9eba9c..d48883993 100644 --- a/discojs/src/default_tasks/index.ts +++ b/discojs/src/default_tasks/index.ts @@ -5,5 +5,5 @@ export { simpleFace } from "./simple_face.js"; export { titanic } from "./titanic.js"; export { wikitext } from "./wikitext.js"; export { tinderDog } from "./tinder_dog.js"; -export { privacyrun } from "./privacyrun.js"; +export { goldfish } from "./goldfish.js"; export { centralizedGPT2FineTune } from "./centralized_gpt2_finetune.js"; From 2dd8187a24e881d8928e08dc8e331f1eab8760a5 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:48:57 +0200 Subject: [PATCH 87/89] fix: don't fail silently --- discojs/src/client/federated/federated_client.ts | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 42b3fea15..0f370b4cd 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -88,7 +88,9 @@ export class FederatedClient extends Client<"federated"> { if (payload != null) { const latestWeights = serialization.weights.decode(payload); model.weights = latestWeights; - } + } else + debug(`[${shortenId(this.ownId)}] received an undefined payload`); + return model; } From d6b60197af2f6c745443ddaab2795e5f949db714 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:53:48 +0200 Subject: [PATCH 88/89] doc: update README cli args --- cli/README.md | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/cli/README.md b/cli/README.md index 650353b0a..fc97ecdcc 100644 --- a/cli/README.md +++ b/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 `./` +- `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, simple-face 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. From b4f341a80a496e651c5c1e1ecd7272d34a3f2a41 Mon Sep 17 00:00:00 2001 From: Julien Vignoud Date: Mon, 31 Aug 2026 16:54:03 +0200 Subject: [PATCH 89/89] fix: rm centralized fine-tuning --- cli/package.json | 1 - cli/src/evaluate_finetuned_gpt2.ts | 294 +++++----- .../evaluate_finetuned_gpt2_full_answer.ts | 531 ------------------ .../centralized_gpt2_finetune.ts | 77 --- 4 files changed, 166 insertions(+), 737 deletions(-) delete mode 100644 cli/src/evaluate_finetuned_gpt2_full_answer.ts delete mode 100644 discojs/src/default_tasks/centralized_gpt2_finetune.ts diff --git a/cli/package.json b/cli/package.json index 118d075c5..7c2f98006 100644 --- a/cli/package.json +++ b/cli/package.json @@ -10,7 +10,6 @@ "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", "build": "tsc --build", "test": ": nothing" diff --git a/cli/src/evaluate_finetuned_gpt2.ts b/cli/src/evaluate_finetuned_gpt2.ts index 4f82e2102..5de9d827d 100644 --- a/cli/src/evaluate_finetuned_gpt2.ts +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -17,7 +17,7 @@ interface Args { } // HOW TO RUN -// npm -w cli run eval_finetuned_gpt2 -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/train_no_exp.txt --maxSamples 100 +// npm -w cli run eval_finetuned_gpt2_full_answer -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/test.txt --maxSamples 100 const PromptFormatNames = [ "answer-colon-space", @@ -30,24 +30,43 @@ type PromptFormatName = (typeof PromptFormatNames)[number]; type PromptFormat = { name: PromptFormatName; makePrompt: (basePrompt: string) => string; - makeContinuation: (option: string) => string; + makeContinuation: (answer: string) => string; +}; + +type Option = { + label: string; + answer: string; +}; + +type ParsedSample = { + basePrompt: string; + answerLabel: string; + answer: string; + options: Option[]; +}; + +type ScoreResult = { + score: number; + promptTokens: number; + continuationTokens: number; + usedInputTokens: number; }; const promptFormats: PromptFormat[] = [ { name: "answer-colon-space", makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, - makeContinuation: (option) => option, + makeContinuation: (answer) => answer, }, { name: "answer-colon", makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, - makeContinuation: (option) => ` ${option}`, + makeContinuation: (answer) => ` ${answer}`, }, { name: "answer-newline", makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, - makeContinuation: (option) => option, + makeContinuation: (answer) => answer, }, ]; @@ -95,23 +114,89 @@ async function loadDataset(filePath: string, limit = -1): Promise { return limit === -1 ? samples : samples.slice(0, limit); } -function parseSample(sample: string) { +function parseAnswerLine( + line: string, +): { label: string; answer?: string } | undefined { + const match = line.trim().match(/^Answer:\s*([A-D])(?:\.\s*(.*))?$/i); + if (match === null) return undefined; + + const label = match[1].toUpperCase(); + const answerText = match[2]?.trim(); + + return { + label, + answer: + answerText === undefined || answerText === "" + ? undefined + : `${label}. ${answerText}`, + }; +} + +function parseOptionLine(line: string): Option | undefined { + const match = line.trim().match(/^([A-D])\.\s*(.+)$/i); + if (match === null) return undefined; + + const label = match[1].toUpperCase(); + return { + label, + answer: `${label}. ${match[2].trim()}`, + }; +} + +function parseSample(sample: string): ParsedSample { const lines = sample.split("\n"); - let answer = ""; + let answerLabel = ""; + let answerFromLine: string | undefined; const promptLines: string[] = []; + const options: Option[] = []; for (const line of lines) { - const trimmed = line.trim(); - if (trimmed.startsWith("Answer:")) { - answer = trimmed.replace("Answer:", "").trim().charAt(0).toUpperCase(); - } else { - promptLines.push(line); + const answer = parseAnswerLine(line); + if (answer !== undefined) { + answerLabel = answer.label; + answerFromLine = answer.answer; + continue; } + + const option = parseOptionLine(line); + if (option !== undefined) { + options.push(option); + } + + promptLines.push(line); + } + + const correctOption = options.find((option) => option.label === answerLabel); + if (correctOption === undefined) { + throw new Error( + `Could not match answer label ${JSON.stringify(answerLabel)} to an option`, + ); + } + + if (answerFromLine !== undefined && answerFromLine !== correctOption.answer) { + throw new Error( + `Answer line ${JSON.stringify(answerFromLine)} does not match option ${JSON.stringify(correctOption.answer)}`, + ); } const basePrompt = promptLines.join("\n").trim(); - return { basePrompt, answer }; + return { + basePrompt, + answerLabel, + answer: correctOption.answer, + options, + }; +} + +function validateOptions(options: Option[], expectedLabels: string[]): boolean { + if (options.length !== expectedLabels.length) return false; + + const labels = options.map((option) => option.label); + return ( + expectedLabels.every((label) => labels.includes(label)) && + new Set(labels).size === expectedLabels.length + ); } async function scoreContinuations( @@ -120,14 +205,7 @@ async function scoreContinuations( prompt: string, continuations: string[], contextLength: number, -): Promise< - { - score: number; - promptTokens: number; - continuationTokens: number; - usedInputTokens: number; - }[] -> { +): Promise { const promptTokens = tokenizer.tokenize(prompt).toArray(); const scoredInputs = continuations.map((continuation) => { const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); @@ -161,55 +239,6 @@ async function scoreContinuations( })); } - const canScoreFromPromptOnly = scoredInputs.every( - ({ continuationStart, continuationTokens }) => - continuationStart === promptTokens.length && continuationTokens === 1, - ); - - if (canScoreFromPromptOnly) { - const offset = Math.max(0, promptTokens.length - contextLength); - const truncatedPromptTokens = promptTokens.slice(offset); - - if (truncatedPromptTokens.length === 0) { - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: truncatedPromptTokens.length, - })); - } - - const inputTensor = tf.tensor2d( - [truncatedPromptTokens], - [1, truncatedPromptTokens.length], - "int32", - ); - - const optionScores = tf.tidy(() => { - const logits = predictTokenLogits(tfModel, inputTensor); - const lastLogits = logits - .slice([0, truncatedPromptTokens.length - 1, 0], [1, 1, -1]) - .reshape([-1]); - const logProbs = tf.logSoftmax(lastLogits); - const continuationTokenIds = scoredInputs.map( - ({ fullTokens, continuationStart }) => fullTokens[continuationStart], - ); - return tf.gather(logProbs, continuationTokenIds); - }); - - const scores = await optionScores.array(); - - inputTensor.dispose(); - optionScores.dispose(); - - return scoredInputs.map((scoredInput, index) => ({ - score: scores[index], - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: truncatedPromptTokens.length, - })); - } - const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ ...truncatedInputTokens, ...Array(maxInputLength - truncatedInputTokens.length).fill(0), @@ -229,8 +258,8 @@ async function scoreContinuations( const { fullTokens, continuationStart, offset, truncatedInputTokens } = scoredInput; - // Usually A/B/C/D is one token and the prompt-only fast path above handles it. - // Keep this fallback for prompt/continuation tokenizer merges and multi-token labels. + // Same ranking as HellaSwag's mean continuation cross-entropy: + // maximize mean log-probability instead of minimizing its negative. for ( let targetPos = continuationStart; targetPos < fullTokens.length; @@ -282,17 +311,15 @@ async function scoreContinuations( scoreCounts[owner]++; }); - const results = scoredInputs.map((scoredInput, index) => { - return { - score: - scoreCounts[index] === 0 - ? Number.NEGATIVE_INFINITY - : scoreSums[index] / scoreCounts[index], - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - }; - }); + const results = scoredInputs.map((scoredInput, index) => ({ + score: + scoreCounts[index] === 0 + ? Number.NEGATIVE_INFINITY + : scoreSums[index] / scoreCounts[index], + promptTokens: promptTokens.length, + continuationTokens: scoredInput.continuationTokens, + usedInputTokens: scoredInput.truncatedInputTokens.length, + })); inputTensor.dispose(); logits.dispose(); @@ -301,7 +328,7 @@ async function scoreContinuations( return results; } -async function benchmarkQA( +async function benchmarkFullAnswers( model: models.GPT, tokenizer: Tokenizer, dataset: string[], @@ -309,26 +336,26 @@ async function benchmarkQA( contextLength: number, savePath?: string, ): Promise { - console.log(`=== QA LOGPROB BENCHMARK (${format.name}) ===`); + console.log(`=== FULL ANSWER LOGPROB BENCHMARK (${format.name}) ===`); console.log(`Context length: ${contextLength}`); const tfModel = model.extract(); let correct = 0; let total = 0; - - const options = ["A", "B", "C", "D"]; - - const confusion: Record> = { - A: { A: 0, B: 0, C: 0, D: 0 }, - B: { A: 0, B: 0, C: 0, D: 0 }, - C: { A: 0, B: 0, C: 0, D: 0 }, - D: { A: 0, B: 0, C: 0, D: 0 }, - }; + const labels = ["A", "B", "C", "D"]; + const confusion: Record> = Object.fromEntries( + labels.map((label) => [ + label, + Object.fromEntries(labels.map((otherLabel) => [otherLabel, 0])), + ]), + ); type PredictionLog = { predicted: string; + predictedAnswer: string; answer: string; + answerText: string; correct: boolean; scores: Record; promptTokens: number; @@ -337,68 +364,75 @@ async function benchmarkQA( }; const logs: PredictionLog[] = []; - const start = Date.now(); for (const sample of dataset) { - const { basePrompt, answer } = parseSample(sample); - - if (!options.includes(answer)) { - console.log("Invalid answer:", JSON.stringify(answer)); + let parsed: ParsedSample; + try { + parsed = parseSample(sample); + } catch (error) { + console.log( + "Invalid sample:", + error instanceof Error ? error.message : error, + ); continue; } - const prompt = format.makePrompt(basePrompt); + if (!validateOptions(parsed.options, labels)) { + console.log( + "Invalid options:", + parsed.options.map((option) => option.label).join(", "), + ); + continue; + } - const scores: number[] = []; - const continuationTokens: number[] = []; - const usedInputTokens: number[] = []; + const prompt = format.makePrompt(parsed.basePrompt); const results = await scoreContinuations( tfModel, tokenizer, prompt, - options.map((opt) => format.makeContinuation(opt)), + parsed.options.map((option) => format.makeContinuation(option.answer)), contextLength, ); - for (const result of results) { - scores.push(result.score); - continuationTokens.push(result.continuationTokens); - usedInputTokens.push(result.usedInputTokens); - } - + const scores = results.map((result) => result.score); let bestIdx = 0; for (let i = 1; i < scores.length; i++) { if (scores[i] > scores[bestIdx]) bestIdx = i; } - const predicted = options[bestIdx]; - - if (predicted === answer) correct++; + const predicted = parsed.options[bestIdx]; + if (predicted.label === parsed.answerLabel) correct++; total++; - if (confusion[answer]?.[predicted] === undefined) { + if (confusion[parsed.answerLabel]?.[predicted.label] === undefined) { throw new Error( - `Unexpected confusion matrix key: answer=${answer}, predicted=${predicted}`, + `Unexpected confusion matrix key: answer=${parsed.answerLabel}, predicted=${predicted.label}`, ); } - confusion[answer][predicted]++; - - const scoreMap = Object.fromEntries( - options.map((opt, i) => [opt, scores[i]]), - ); + confusion[parsed.answerLabel][predicted.label]++; logs.push({ - predicted, - answer, - correct: predicted === answer, - scores: scoreMap, + predicted: predicted.label, + predictedAnswer: predicted.answer, + answer: parsed.answerLabel, + answerText: parsed.answer, + correct: predicted.label === parsed.answerLabel, + scores: Object.fromEntries( + parsed.options.map((option, i) => [option.label, scores[i]]), + ), promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, continuationTokens: Object.fromEntries( - options.map((opt, i) => [opt, continuationTokens[i]]), + parsed.options.map((option, i) => [ + option.label, + results[i].continuationTokens, + ]), ), usedInputTokens: Object.fromEntries( - options.map((opt, i) => [opt, usedInputTokens[i]]), + parsed.options.map((option, i) => [ + option.label, + results[i].usedInputTokens, + ]), ), }); @@ -407,6 +441,10 @@ async function benchmarkQA( } } + if (total === 0) { + throw new Error("No valid samples were evaluated"); + } + const accuracy = correct / total; const duration = ((Date.now() - start) / 1000).toFixed(2); @@ -419,7 +457,7 @@ async function benchmarkQA( console.table(confusion); console.log("\nPer-class accuracy:"); - for (const cls of options) { + for (const cls of labels) { const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); const correctCls = confusion[cls][cls]; const acc = totalCls ? (correctCls / totalCls) * 100 : 0; @@ -477,7 +515,7 @@ async function main() { ? args.savePath : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); - await benchmarkQA( + await benchmarkFullAnswers( model, tokenizer, dataset, diff --git a/cli/src/evaluate_finetuned_gpt2_full_answer.ts b/cli/src/evaluate_finetuned_gpt2_full_answer.ts deleted file mode 100644 index 5de9d827d..000000000 --- a/cli/src/evaluate_finetuned_gpt2_full_answer.ts +++ /dev/null @@ -1,531 +0,0 @@ -import "@tensorflow/tfjs-node"; -import * as tf from "@tensorflow/tfjs"; -import fs from "node:fs/promises"; -import { parse } from "ts-command-line-args"; -import { models, Tokenizer } from "@epfml/discojs"; -import { loadModelFromDisk } from "@epfml/discojs-node"; - -interface Args { - modelPath: string; - testPath: string; - maxSamples?: number; - savePath?: string; - compareFormats?: boolean; - promptFormat?: PromptFormatName; - contextLength?: number; - help?: boolean; -} - -// HOW TO RUN -// npm -w cli run eval_finetuned_gpt2_full_answer -- --modelPath absolute_path_to_model/model.json --testPath absolute_path_to_test_data/test.txt --maxSamples 100 - -const PromptFormatNames = [ - "answer-colon-space", - "answer-colon", - "answer-newline", -] as const; - -type PromptFormatName = (typeof PromptFormatNames)[number]; - -type PromptFormat = { - name: PromptFormatName; - makePrompt: (basePrompt: string) => string; - makeContinuation: (answer: string) => string; -}; - -type Option = { - label: string; - answer: string; -}; - -type ParsedSample = { - basePrompt: string; - answerLabel: string; - answer: string; - options: Option[]; -}; - -type ScoreResult = { - score: number; - promptTokens: number; - continuationTokens: number; - usedInputTokens: number; -}; - -const promptFormats: PromptFormat[] = [ - { - name: "answer-colon-space", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer: `, - makeContinuation: (answer) => answer, - }, - { - name: "answer-colon", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer:`, - makeContinuation: (answer) => ` ${answer}`, - }, - { - name: "answer-newline", - makePrompt: (basePrompt) => `${basePrompt}\nAnswer:\n`, - makeContinuation: (answer) => answer, - }, -]; - -function castPromptFormatName(raw: string): PromptFormatName { - for (const name of PromptFormatNames) { - if (raw === name) return name; - } - throw new Error(`Invalid promptFormat: ${raw}`); -} - -function commonPrefixLength(left: number[], right: number[]): number { - const maxLength = Math.min(left.length, right.length); - - for (let i = 0; i < maxLength; i++) { - if (left[i] !== right[i]) return i; - } - - return maxLength; -} - -function predictTokenLogits( - tfModel: tf.LayersModel, - inputTensor: tf.Tensor2D, -): tf.Tensor3D { - const logits = tfModel.predict(inputTensor); - if (Array.isArray(logits)) { - throw new Error("Expected GPT model to return a single logits tensor"); - } - if (logits.rank !== 3) { - logits.dispose(); - throw new Error( - `Expected GPT logits to have rank 3, got rank ${logits.rank}`, - ); - } - return logits as tf.Tensor3D; -} - -async function loadDataset(filePath: string, limit = -1): Promise { - const text = await fs.readFile(filePath, "utf-8"); - const samples = text - .split("<|endoftext|>") - .map((sample) => sample.replaceAll("<|startoftext|>", "").trim()) - .filter((sample) => sample !== ""); - - return limit === -1 ? samples : samples.slice(0, limit); -} - -function parseAnswerLine( - line: string, -): { label: string; answer?: string } | undefined { - const match = line.trim().match(/^Answer:\s*([A-D])(?:\.\s*(.*))?$/i); - if (match === null) return undefined; - - const label = match[1].toUpperCase(); - const answerText = match[2]?.trim(); - - return { - label, - answer: - answerText === undefined || answerText === "" - ? undefined - : `${label}. ${answerText}`, - }; -} - -function parseOptionLine(line: string): Option | undefined { - const match = line.trim().match(/^([A-D])\.\s*(.+)$/i); - if (match === null) return undefined; - - const label = match[1].toUpperCase(); - return { - label, - answer: `${label}. ${match[2].trim()}`, - }; -} - -function parseSample(sample: string): ParsedSample { - const lines = sample.split("\n"); - - let answerLabel = ""; - let answerFromLine: string | undefined; - const promptLines: string[] = []; - const options: Option[] = []; - - for (const line of lines) { - const answer = parseAnswerLine(line); - if (answer !== undefined) { - answerLabel = answer.label; - answerFromLine = answer.answer; - continue; - } - - const option = parseOptionLine(line); - if (option !== undefined) { - options.push(option); - } - - promptLines.push(line); - } - - const correctOption = options.find((option) => option.label === answerLabel); - if (correctOption === undefined) { - throw new Error( - `Could not match answer label ${JSON.stringify(answerLabel)} to an option`, - ); - } - - if (answerFromLine !== undefined && answerFromLine !== correctOption.answer) { - throw new Error( - `Answer line ${JSON.stringify(answerFromLine)} does not match option ${JSON.stringify(correctOption.answer)}`, - ); - } - - const basePrompt = promptLines.join("\n").trim(); - return { - basePrompt, - answerLabel, - answer: correctOption.answer, - options, - }; -} - -function validateOptions(options: Option[], expectedLabels: string[]): boolean { - if (options.length !== expectedLabels.length) return false; - - const labels = options.map((option) => option.label); - return ( - expectedLabels.every((label) => labels.includes(label)) && - new Set(labels).size === expectedLabels.length - ); -} - -async function scoreContinuations( - tfModel: tf.LayersModel, - tokenizer: Tokenizer, - prompt: string, - continuations: string[], - contextLength: number, -): Promise { - const promptTokens = tokenizer.tokenize(prompt).toArray(); - const scoredInputs = continuations.map((continuation) => { - const fullTokens = tokenizer.tokenize(prompt + continuation).toArray(); - const continuationStart = commonPrefixLength(promptTokens, fullTokens); - const continuationTokens = fullTokens.length - continuationStart; - const inputTokens = fullTokens.slice(0, -1); - const offset = Math.max(0, inputTokens.length - contextLength); - const truncatedInputTokens = inputTokens.slice(offset); - - return { - fullTokens, - continuationStart, - continuationTokens, - offset, - truncatedInputTokens, - }; - }); - - const maxInputLength = scoredInputs.reduce( - (maxLength, scoredInput) => - Math.max(maxLength, scoredInput.truncatedInputTokens.length), - 0, - ); - - if (maxInputLength === 0) { - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); - } - - const paddedInputs = scoredInputs.map(({ truncatedInputTokens }) => [ - ...truncatedInputTokens, - ...Array(maxInputLength - truncatedInputTokens.length).fill(0), - ]); - - const inputTensor = tf.tensor2d( - paddedInputs, - [paddedInputs.length, maxInputLength], - "int32", - ); - - const targetIndexes: number[][] = []; - const targetTokenIds: number[] = []; - const targetOwners: number[] = []; - - scoredInputs.forEach((scoredInput, batchIdx) => { - const { fullTokens, continuationStart, offset, truncatedInputTokens } = - scoredInput; - - // Same ranking as HellaSwag's mean continuation cross-entropy: - // maximize mean log-probability instead of minimizing its negative. - for ( - let targetPos = continuationStart; - targetPos < fullTokens.length; - targetPos++ - ) { - const targetToken = fullTokens[targetPos]; - const logitPos = targetPos - 1 - offset; - if (logitPos < 0 || logitPos >= truncatedInputTokens.length) continue; - targetIndexes.push([batchIdx, logitPos]); - targetTokenIds.push(targetToken); - targetOwners.push(batchIdx); - } - }); - - if (targetIndexes.length === 0) { - inputTensor.dispose(); - return scoredInputs.map((scoredInput) => ({ - score: Number.NEGATIVE_INFINITY, - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); - } - - const logits = predictTokenLogits(tfModel, inputTensor); - const targetLogProbs = tf.tidy(() => { - const targetIndexTensor = tf.tensor2d( - targetIndexes, - [targetIndexes.length, 2], - "int32", - ); - const targetTokenIndexTensor = tf.tensor2d( - targetTokenIds.map((targetTokenId, index) => [index, targetTokenId]), - [targetTokenIds.length, 2], - "int32", - ); - const targetLogits = tf.gatherND(logits, targetIndexTensor) as tf.Tensor2D; - const logProbs = tf.logSoftmax(targetLogits, -1); - return tf.gatherND(logProbs, targetTokenIndexTensor); - }); - - const targetScores = (await targetLogProbs.array()) as number[]; - const scoreSums = Array(scoredInputs.length).fill(0) as number[]; - const scoreCounts = Array(scoredInputs.length).fill(0) as number[]; - - targetScores.forEach((score, index) => { - const owner = targetOwners[index]; - scoreSums[owner] += score; - scoreCounts[owner]++; - }); - - const results = scoredInputs.map((scoredInput, index) => ({ - score: - scoreCounts[index] === 0 - ? Number.NEGATIVE_INFINITY - : scoreSums[index] / scoreCounts[index], - promptTokens: promptTokens.length, - continuationTokens: scoredInput.continuationTokens, - usedInputTokens: scoredInput.truncatedInputTokens.length, - })); - - inputTensor.dispose(); - logits.dispose(); - targetLogProbs.dispose(); - - return results; -} - -async function benchmarkFullAnswers( - model: models.GPT, - tokenizer: Tokenizer, - dataset: string[], - format: PromptFormat, - contextLength: number, - savePath?: string, -): Promise { - console.log(`=== FULL ANSWER LOGPROB BENCHMARK (${format.name}) ===`); - console.log(`Context length: ${contextLength}`); - - const tfModel = model.extract(); - - let correct = 0; - let total = 0; - const labels = ["A", "B", "C", "D"]; - const confusion: Record> = Object.fromEntries( - labels.map((label) => [ - label, - Object.fromEntries(labels.map((otherLabel) => [otherLabel, 0])), - ]), - ); - - type PredictionLog = { - predicted: string; - predictedAnswer: string; - answer: string; - answerText: string; - correct: boolean; - scores: Record; - promptTokens: number; - continuationTokens: Record; - usedInputTokens: Record; - }; - - const logs: PredictionLog[] = []; - const start = Date.now(); - - for (const sample of dataset) { - let parsed: ParsedSample; - try { - parsed = parseSample(sample); - } catch (error) { - console.log( - "Invalid sample:", - error instanceof Error ? error.message : error, - ); - continue; - } - - if (!validateOptions(parsed.options, labels)) { - console.log( - "Invalid options:", - parsed.options.map((option) => option.label).join(", "), - ); - continue; - } - - const prompt = format.makePrompt(parsed.basePrompt); - const results = await scoreContinuations( - tfModel, - tokenizer, - prompt, - parsed.options.map((option) => format.makeContinuation(option.answer)), - contextLength, - ); - - const scores = results.map((result) => result.score); - let bestIdx = 0; - for (let i = 1; i < scores.length; i++) { - if (scores[i] > scores[bestIdx]) bestIdx = i; - } - - const predicted = parsed.options[bestIdx]; - if (predicted.label === parsed.answerLabel) correct++; - total++; - - if (confusion[parsed.answerLabel]?.[predicted.label] === undefined) { - throw new Error( - `Unexpected confusion matrix key: answer=${parsed.answerLabel}, predicted=${predicted.label}`, - ); - } - confusion[parsed.answerLabel][predicted.label]++; - - logs.push({ - predicted: predicted.label, - predictedAnswer: predicted.answer, - answer: parsed.answerLabel, - answerText: parsed.answer, - correct: predicted.label === parsed.answerLabel, - scores: Object.fromEntries( - parsed.options.map((option, i) => [option.label, scores[i]]), - ), - promptTokens: results[0]?.promptTokens ?? tokenizer.tokenize(prompt).size, - continuationTokens: Object.fromEntries( - parsed.options.map((option, i) => [ - option.label, - results[i].continuationTokens, - ]), - ), - usedInputTokens: Object.fromEntries( - parsed.options.map((option, i) => [ - option.label, - results[i].usedInputTokens, - ]), - ), - }); - - if (total % 50 === 0) { - console.log(`Processed ${total} samples...`); - } - } - - if (total === 0) { - throw new Error("No valid samples were evaluated"); - } - - const accuracy = correct / total; - const duration = ((Date.now() - start) / 1000).toFixed(2); - - console.log("\n========================="); - console.log(`Accuracy: ${(accuracy * 100).toFixed(2)}%`); - console.log(`Time: ${duration}s`); - console.log("=========================\n"); - - console.log("Confusion Matrix:"); - console.table(confusion); - - console.log("\nPer-class accuracy:"); - for (const cls of labels) { - const totalCls = Object.values(confusion[cls]).reduce((a, b) => a + b, 0); - const correctCls = confusion[cls][cls]; - const acc = totalCls ? (correctCls / totalCls) * 100 : 0; - - console.log(`${cls}: ${acc.toFixed(2)}%`); - } - - if (savePath) { - await fs.writeFile(savePath, JSON.stringify(logs, null, 2)); - console.log(`Saved results to ${savePath}`); - } - - return accuracy; -} - -async function main() { - const args = parse({ - modelPath: { type: String }, - testPath: { type: String }, - maxSamples: { type: Number, optional: true, defaultValue: 100 }, - savePath: { type: String, optional: true }, - compareFormats: { type: Boolean, optional: true, defaultValue: false }, - promptFormat: { - type: (raw: string) => castPromptFormatName(raw), - optional: true, - defaultValue: "answer-colon-space", - }, - contextLength: { type: Number, optional: true }, - help: { type: Boolean, optional: true }, - }); - - console.log("Loading tokenizer..."); - const tokenizer = await Tokenizer.from_pretrained("Xenova/gpt2"); - - console.log("Loading model..."); - const model = await loadModelFromDisk(args.modelPath); - - if (!(model instanceof models.GPT)) { - throw new Error("Model must be GPT"); - } - - console.log("Loading dataset..."); - const dataset = await loadDataset(args.testPath, args.maxSamples); - - console.log(`Loaded ${dataset.length} samples`); - - const contextLength = args.contextLength ?? model.config.contextLength; - const formats = args.compareFormats - ? promptFormats - : promptFormats.filter((format) => format.name === args.promptFormat); - - for (const format of formats) { - const savePath = - args.savePath === undefined || formats.length === 1 - ? args.savePath - : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); - - await benchmarkFullAnswers( - model, - tokenizer, - dataset, - format, - contextLength, - savePath, - ); - } - - console.log("Done."); -} - -main().catch(console.error); diff --git a/discojs/src/default_tasks/centralized_gpt2_finetune.ts b/discojs/src/default_tasks/centralized_gpt2_finetune.ts deleted file mode 100644 index 4c0e60508..000000000 --- a/discojs/src/default_tasks/centralized_gpt2_finetune.ts +++ /dev/null @@ -1,77 +0,0 @@ -import type { TaskProvider } from "../index.js"; -import { Tokenizer, models, serialization } from "../index.js"; - -export const centralizedGPT2FineTune: TaskProvider<"text", "local"> = { - async getTask() { - return { - id: "centralized-gpt2-finetune", - dataType: "text", - displayInformation: { - title: "GPT Privacy-Preserving Fine-tuning", - summary: { - preview: - "Fine-tune a pre-trained GPT model collaboratively and privately.", - overview: - "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning.", - }, - model: [ - "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", - "The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer.", - "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", - "Context length is kept at 512 to match the pre-trained model, with batch size at 8.", - ].join(" "), - dataFormatInformation: - "You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.", - dataExample: - "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", - }, - trainingInformation: { - scheme: "local", - aggregationStrategy: "mean", - epochs: 1, - validationSplit: 0.1, - roundDuration: 1, - // Last context segment may be shorter than context length, so it will be dropped (TODO: implement padding to avoid this) - batchSize: 8, - tokenizer: await Tokenizer.from_pretrained("Xenova/gpt2"), - contextLength: 512, - tensorBackend: "gpt", - }, - }; - }, - - async getModel() { - // Load the pre-trained ONNX-converted model from Google Cloud Storage - // The model should be in DiscoJS serialization format (created by onnx-converter) - - const modelUrl = - "https://storage.googleapis.com/deai-313515.appspot.com/model_ctx_512.json"; - - try { - const response = await fetch(modelUrl); - if (!response.ok) { - throw new Error(`HTTP ${response.status}: ${response.statusText}`); - } - - const arrayBuffer = await response.arrayBuffer(); - const encodedData = new Uint8Array(arrayBuffer); - - const model = await serialization.model.decode(encodedData); - - if (!(model instanceof models.GPT)) { - throw new Error("Loaded model is not a GPT model"); - } - - console.log( - "Successfully loaded pre-trained GPT model from Google Cloud Storage", - ); - - return model; - } catch (error) { - console.error("Failed to load model from Google Cloud Storage:", error); - throw new Error( - `Could not load model from ${modelUrl}. Make sure the URL is correct and the model exists in DiscoJS serialization format.`, - ); - } - }, -};