From c21d55bd1b202d47df5713532642323ccb946fcb Mon Sep 17 00:00:00 2001 From: mina5rovic Date: Thu, 16 Apr 2026 13:49:48 +0200 Subject: [PATCH 01/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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/59] 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)