diff --git a/.knip.json b/.knip.json index 8336807d9..5c8edfad6 100644 --- a/.knip.json +++ b/.knip.json @@ -12,7 +12,10 @@ "entry": [ "src/benchmark_gpt.ts", "src/hellaswag_gpt.ts", - "src/train_gpt.ts" + "src/train_gpt.ts", + "src/evaluate_finetuned_gpt2_full_answer.ts", + "src/evaluate_finetuned_gpt2.ts", + "src/measure_memorization_gpt2.ts" ] }, "onnx-converter": { diff --git a/Dockerfile b/Dockerfile index 9c9631f70..16492951e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,7 @@ # TODO freeze to 22 until tfjs#8425 is merged FROM node:22 AS builder -RUN npm i -g pnpm +RUN npm i -g pnpm@11.11.0 WORKDIR /disco diff --git a/cli/README.md b/cli/README.md index daffbd266..fc97ecdcc 100644 --- a/cli/README.md +++ b/cli/README.md @@ -36,14 +36,34 @@ Non-mandatory fields will automatically use values from the task specification. - `testID`: (mandatory) arbitrary test ID defined by the user for the test run - `task`: (mandatory) pre-defined task (adding a new task is described in the next section) - `numberOfUsers`: number of users participating in the learning round -- `save`: whether to save the logs of the test run +- `host`: URL of the server to connect to, defaults to `http://localhost:8080` +- `outputPath`: path to save logs and models, defaults to `./` +- `saveLogs`: whether to save the logs of the test run +- `saveModel`: whether to save the trained model to disk +- `saveCheckpoints`: whether to save each client model after every completed round/aggregation + +### Dataset arguments + +- `datasetPath`: path to the training dataset +- `validationDatasetPath`: path to a separate validation dataset shared by all clients, takes precedence over `validationSplit` ### Learning hyperparameters - `epochs`: total number of training epochs -- `roundDuration`: number of epochs per round +- `roundDuration`: number of epochs per round, ignored if `roundIterations` is set. +- `roundIterations`: number of iterations per round, takes precedence over `roundDuration` - `batchSize`: batch size -- `validationSplit`: ratio of the validation set used for evaluation +- `validationSplit`: fraction of each client's training data used for validation, ignored when `validationDatasetPath` is set; 0 disables split-based validation +- `validationFrequency`: how often to validate. Validate the first aggregation round and every N rounds after it; defaults to every round, 0 disables validation metrics +- `validationMode`: when to run the validation: `before` model aggregation (default), `after`, or `both` +- `learningRate`: override the learning rate (GPT text tasks only) + +### Goldfish loss parameters (GPT text tasks only) + +- `goldfishLoss`: train with the [goldfish loss](https://arxiv.org/abs/2406.10209), which drops a subset of target tokens from the loss to mitigate memorization +- `goldfishK`: drop modulus k, a target token is dropped if hash(context) mod k == 0 +- `goldfishH`: localized hash context length +- `goldfishPadTokenId`: (optional) padding token id to exclude from the goldfish loss denominator ### Aggregator parameters @@ -57,7 +77,7 @@ Non-mandatory fields will automatically use values from the task specification. ## Adding new tasks -The CLI can be used on several pre-defined tasks: titanic, lus_covid and CIFAR10. In order +The CLI can be used on several pre-defined tasks: `cifar10`, `lus_covid`, `mnist`, `simple_face`, `tinder_dog`, `titanic` and `goldfish` (GPT-2 fine-tuning). In order to understand how to add a new task have a look at [TASK.md](../docs/TASK.md). Once a new task has been defined in `discojs`, it can be loaded in [data.ts](./src/data.ts) as it is already implemented for current tasks. There are currently [multiple classes](../discojs-node/src/loaders) you can use to load data using Node.js and preprocess data: loadImagesInDir, loadCSV and loadText. diff --git a/cli/package.json b/cli/package.json index c8f781c1e..7c2f98006 100644 --- a/cli/package.json +++ b/cli/package.json @@ -9,6 +9,8 @@ "benchmark_gpt": "pnpm run build && node dist/benchmark_gpt.js", "train_gpt": "pnpm run build && node dist/train_gpt.js", "hellaswag_gpt": "pnpm run build && node dist/hellaswag_gpt.js", + "eval_finetuned_gpt2": "pnpm run build && node dist/evaluate_finetuned_gpt2.js", + "measure_memorization_gpt2": "pnpm run build && node dist/measure_memorization_gpt2.js", "build": "tsc --build", "test": ": nothing" }, @@ -17,7 +19,9 @@ "dependencies": { "@epfml/discojs": "workspace:", "@epfml/discojs-node": "workspace:", + "@tensorflow/tfjs": "catalog:", "@tensorflow/tfjs-node": "catalog:", + "debug": "catalog:", "immutable": "catalog:", "server": "workspace:" }, diff --git a/cli/src/args.ts b/cli/src/args.ts index b1412ebdf..a06cdcaec 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, GPT } from "@epfml/discojs"; type AggregationStrategy = "mean" | "byzantine" | "secure"; @@ -11,14 +11,56 @@ function parseAggregator(raw: string): AggregationStrategy { else throw new Error(`Aggregator ${raw} is not supported.`); } +type ValidationMode = "before" | "after" | "both"; + +function parseValidationMode(raw: string): ValidationMode { + if (raw === "before" || raw === "after" || raw === "both") return raw; + else + throw new Error( + `Validation mode ${raw} is not supported, expected "before", "after" or "both".`, + ); +} + export interface BenchmarkArguments { provider: TaskProvider; testID: string; numberOfUsers: number; epochs: number; roundDuration: number; + roundIterations?: number; batchSize: number; + /** + * Fraction of each client's training dataset reserved for validation. + * Ignored when `validationDatasetPath` is set. A value of 0 leaves the + * client without validation data unless `validationDatasetPath` is set. + */ validationSplit: number; + /** + * Validate the first aggregation round and every N rounds after it. If + * omitted, validation runs every round; 0 disables validation metrics. This + * only controls when validation runs, not whether its data comes from + * `validationSplit` or `validationDatasetPath`. + */ + validationFrequency?: number; + /** + * When to run validation relative to weight aggregation: "before" (on the + * local model), "after" (on the freshly aggregated global model), or "both". + * Defaults to "before". + */ + validationMode?: ValidationMode; + datasetPath?: string; + /** + * Path to a separate validation dataset. When set, this dataset is shared + * by all clients and takes precedence over `validationSplit`, including + * when `validationSplit` is non-zero. + */ + validationDatasetPath?: string; + outputPath?: string; + goldfishLoss: boolean; + goldfishK: number; + goldfishH: number; + goldfishPadTokenId?: number; + learningRate?: number; // DP epsilon?: number; @@ -33,12 +75,16 @@ export interface BenchmarkArguments { // Secure aggregator maxShareValue?: number; - save: boolean; + saveLogs: boolean; + saveModel: boolean; + saveCheckpoints: boolean; host: URL; } type BenchmarkUnsafeArguments = Omit & { task: string; + datasetPath?: string; + validationDatasetPath?: string; help?: boolean; }; @@ -46,11 +92,7 @@ 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", - }, + testID: { type: String, alias: "i", description: "ID of the testcase" }, task: { type: String, alias: "t", @@ -75,6 +117,12 @@ const unsafeArgs = parse( description: "Round duration (in epochs)", defaultValue: 2, }, + roundIterations: { + type: Number, + description: + "For GPT text tasks, aggregate every N training batches without rewinding the dataset", + optional: true, + }, batchSize: { type: Number, alias: "b", @@ -84,15 +132,87 @@ const unsafeArgs = parse( validationSplit: { type: Number, alias: "v", - description: "Validation dataset ratio", + description: + "Fraction of each client's training data used for validation. Ignored when --validationDatasetPath is set; 0 disables split-based validation.", defaultValue: 0.2, }, - save: { + validationFrequency: { + type: Number, + description: + "Validate the first aggregation round and every N rounds after it. Defaults to every round; use 0 to disable validation metrics.", + optional: true, + }, + validationMode: { + type: parseValidationMode, + typeLabel: "before|after|both", + description: + "When to run validation relative to weight aggregation: before (local model), after (aggregated global model), or both. Defaults to before.", + optional: true, + }, + datasetPath: { + type: String, + alias: "d", + description: "Path to the dataset", + optional: true, + }, + validationDatasetPath: { + type: String, + alias: "V", + description: + "Path to a separate validation dataset shared by all clients. Takes precedence over --validationSplit.", + optional: true, + }, + outputPath: { + type: String, + alias: "o", + description: "Path to save logs and models. Defaults to ./", + optional: true, + }, + goldfishLoss: { + type: Boolean, + description: "Use Goldfish loss for GPT text tasks", + defaultValue: false, + }, + goldfishK: { + type: Number, + description: + "Goldfish loss drop modulus k. Drops target if hash(context) mod k == 0", + defaultValue: 4, + }, + goldfishH: { + type: Number, + description: "Goldfish loss localized hash context length", + defaultValue: 13, + }, + goldfishPadTokenId: { + type: Number, + description: + "Optional padding token id to exclude from Goldfish loss denominator", + optional: true, + }, + learningRate: { + type: Number, + description: "Override learning rate for GPT text tasks", + optional: true, + }, + saveLogs: { type: Boolean, alias: "s", description: "Save logs of benchmark", defaultValue: false, }, + saveModel: { + type: Boolean, + alias: "m", + description: "Save trained model to disk", + defaultValue: false, + }, + saveCheckpoints: { + type: Boolean, + description: + "Save each client model after every completed round/aggregation", + defaultValue: false, + }, host: { type: (raw: string) => new URL(raw), typeLabel: "URL", @@ -172,17 +292,18 @@ const unsafeArgs = parse( const supportedTasks = Map( await Promise.all( - Set.of>( + Set.of>( defaultTasks.cifar10, defaultTasks.lusCovid, defaultTasks.titanic, defaultTasks.tinderDog, defaultTasks.mnist, + defaultTasks.goldfish, ).map( async (t) => [(await t.getTask()).id, t] as [ string, - TaskProvider<"image" | "tabular", Network>, + TaskProvider<"image" | "tabular" | "text", Network>, ], ), ), @@ -204,6 +325,46 @@ export const args: BenchmarkArguments = { task.trainingInformation.roundDuration = unsafeArgs.roundDuration; task.trainingInformation.epochs = unsafeArgs.epochs; task.trainingInformation.validationSplit = unsafeArgs.validationSplit; + task.trainingInformation.roundIterations = unsafeArgs.roundIterations; + task.trainingInformation.validationFrequency = + unsafeArgs.validationFrequency; + task.trainingInformation.validationMode = unsafeArgs.validationMode; + + if (unsafeArgs.goldfishLoss) { + if ( + task.dataType !== "text" || + task.trainingInformation.tensorBackend !== "gpt" + ) + throw new Error("Goldfish loss is only supported for GPT text tasks"); + if (!Number.isInteger(unsafeArgs.goldfishK) || unsafeArgs.goldfishK < 1) + throw new Error("goldfishK must be a positive integer"); + if (!Number.isInteger(unsafeArgs.goldfishH) || unsafeArgs.goldfishH < 1) + throw new Error("goldfishH must be a positive integer"); + + task.trainingInformation.goldfishLoss = { + enabled: true, + k: unsafeArgs.goldfishK, + h: unsafeArgs.goldfishH, + padTokenId: unsafeArgs.goldfishPadTokenId, + }; + } + + if (unsafeArgs.learningRate !== undefined) { + if ( + task.dataType !== "text" || + task.trainingInformation.tensorBackend !== "gpt" + ) + throw new Error( + "learningRate override is only supported for GPT text tasks", + ); + if ( + !Number.isFinite(unsafeArgs.learningRate) || + unsafeArgs.learningRate <= 0 + ) + throw new Error("learningRate must be a positive finite number"); + + task.trainingInformation.learningRate = unsafeArgs.learningRate; + } const { aggregator, clippingRadius, maxIterations, beta, maxShareValue } = unsafeArgs; @@ -277,6 +438,30 @@ export const args: BenchmarkArguments = { return task; }, - modelCard: provider.modelCard, + modelCard: { + card: provider.modelCard.card, + async getModel() { + const model = await provider.modelCard.getModel(); + + if (unsafeArgs.learningRate !== undefined) { + if (!(model instanceof GPT)) + throw new Error( + "learningRate override is only supported for GPT models", + ); + if ( + !Number.isFinite(unsafeArgs.learningRate) || + unsafeArgs.learningRate <= 0 + ) + throw new Error("learningRate must be a positive finite number"); + + model.setLearningRate(unsafeArgs.learningRate); + console.log( + `Overriding GPT learning rate to ${unsafeArgs.learningRate}`, + ); + } + + return model; + }, + }, }, }; diff --git a/cli/src/cli.ts b/cli/src/cli.ts index 8ae0b7201..9193e98af 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, @@ -14,28 +14,85 @@ import type { Task, TaskProvider, Network, + Model, } from "@epfml/discojs"; import { Disco, getAggregator, getClient } from "@epfml/discojs"; +import { 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"); + +let checkpointQueue = Promise.resolve(); + +function getOutputDir(): string { + return args.outputPath ?? path.join(".", `${args.testID}`); +} + +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; + } + gc(); +} + +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 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, url: URL, data: Dataset, + validationData: Dataset | undefined, userIndex: number, numberOfUsers: number, ): Promise> { - // cast as typescript isn't good with generics + debug(`Starting runUser for client ${userIndex}`); const trainingScheme = task.trainingInformation.scheme as N; const aggregator = getAggregator(task); const client = getClient(trainingScheme, url, task, aggregator); - const disco = new Disco(task, client, { scheme: trainingScheme }); + const disco = new Disco(task, client, { + scheme: trainingScheme, + preprocessOnce: false, + debugLabel: `client${userIndex}`, + }); - 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`); @@ -43,30 +100,58 @@ 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" }); } try { - for await (const log of disco.trainSummary(data)) { + 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 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}`, + ); + } // saving the entire per-user logs - if (args.save) { + 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, - client.ownId, + clientId, finalLog, ); @@ -115,23 +200,39 @@ async function main( const dataSplits = await Promise.all( Range(0, numberOfUsers).map(async (i) => - getTaskData(task.id, i, numberOfUsers), + 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, + ) + ).cached() as Dataset; + } + const logs = await Promise.all( dataSplits.map((data, i) => runUser( task, 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/cli/src/data.ts b/cli/src/data.ts index ce4ee72b3..7735203dd 100644 --- a/cli/src/data.ts +++ b/cli/src/data.ts @@ -1,10 +1,60 @@ import path from "node:path"; -import type { Dataset } from "@epfml/discojs"; -import { extractColumn } from "@epfml/discojs"; +import { createReadStream } from "node:fs"; +import { Dataset, extractColumn } from "@epfml/discojs"; import type { DataFormat, DataType, Image, Task } from "@epfml/discojs"; import { loadCSV, loadImage, loadImagesInDir } from "@epfml/discojs-node"; import { Repeat } from "immutable"; +function loadTextSamples( + 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(); + const shouldYield = + userIdx === undefined || + totalClient === undefined || + sampleIndex % totalClient === userIdx; + + if (sample !== "" && shouldYield) { + yield sample; + } + + sampleIndex++; + buffer = buffer.slice(delimiterIndex + sampleDelimiter.length); + delimiterIndex = buffer.indexOf(sampleDelimiter); + } + } + + const trailingSample = buffer.trim(); + const shouldYieldTrailing = + userIdx === undefined || + totalClient === undefined || + sampleIndex % totalClient === userIdx; + + if (trailingSample !== "" && shouldYieldTrailing) { + yield trailingSample; + } + }); +} + async function loadLusCovidData( userIdx: number, totalClient: number, @@ -76,7 +126,15 @@ export async function getTaskData( taskID: Task.ID, userIdx: number, totalClient: number, + datasetPath?: string, + isValidation?: boolean, + validationDatasetPath?: string, ): Promise> { + if (validationDatasetPath && taskID !== "goldfish") + throw new Error( + "validationDatasetPath is currently only supported for the goldfish task", + ); + switch (taskID) { case "titanic": case "titanic_decentralized": @@ -99,6 +157,21 @@ export async function getTaskData( case "mnist_federated": case "mnist": return loadData("mnist", userIdx) as Dataset; + case "goldfish": { + 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 loadTextSamples(filePath) 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 new file mode 100644 index 000000000..b3409c657 --- /dev/null +++ b/cli/src/evaluate_finetuned_gpt2.ts @@ -0,0 +1,531 @@ +import "@tensorflow/tfjs-node"; +import * as tf from "@tensorflow/tfjs"; +import fs from "node:fs/promises"; +import { parse } from "ts-command-line-args"; +import { GPT, 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: 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 GPT)) { + throw new Error("Model must be GPT"); + } + + console.log("Loading dataset..."); + const dataset = await loadDataset(args.testPath, args.maxSamples); + + console.log(`Loaded ${dataset.length} samples`); + + const contextLength = args.contextLength ?? model.config.contextLength; + const formats = args.compareFormats + ? promptFormats + : promptFormats.filter((format) => format.name === args.promptFormat); + + for (const format of formats) { + const savePath = + args.savePath === undefined || formats.length === 1 + ? args.savePath + : args.savePath.replace(/(\.[^.]+)?$/, `.${format.name}$1`); + + await benchmarkFullAnswers( + model, + tokenizer, + dataset, + format, + contextLength, + savePath, + ); + } + + console.log("Done."); +} + +main().catch(console.error); diff --git a/cli/src/measure_memorization_gpt2.ts b/cli/src/measure_memorization_gpt2.ts new file mode 100644 index 000000000..d8ba7aa53 --- /dev/null +++ b/cli/src/measure_memorization_gpt2.ts @@ -0,0 +1,566 @@ +import "@tensorflow/tfjs-node"; +import * as tf from "@tensorflow/tfjs"; +import fs from "node:fs/promises"; +import { parse } from "ts-command-line-args"; + +import { GPT, Tokenizer } from "@epfml/discojs"; +import { loadModelFromDisk } from "@epfml/discojs-node"; + +interface Args { + modelPath: string; + dataPath: string; + maxRecords: number; + 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; + promptLength: number; + splitIndex: number; + exactMatch: boolean; + bleu: number; + memorizedByBleu: boolean; + promptText: string; + referenceText: string; + generatedText: string; +}; + +type TokenLengthStats = { + min: number; + p50: number; + p90: number; + max: number; + average: number; +}; + +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 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 () => { + 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 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; + + 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 sampleGenerateGPT2( + model: GPT, + inputIds: number[], + maxNewTokens: number, + maxContextLength: number, + decodingStrategy: DecodingStrategy, + temperature: number, + topK: number, + seed: 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]; + } + return output; + }); + + 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); + + 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(); + 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, + }, + 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", + }, + }, + { + helpArg: "help", + headerContentSections: [ + { + header: "GPT-2 Unintended Memorization", + content: + "Measures extractable memorization via greedy suffix generation.", + }, + ], + }, + ); + + 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 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 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`); + + 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 < 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" + }`, + ); + } + + 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, + maxEligiblePromptLength, + ids.length - args.suffixLength, + ); + const reference = ids.slice(splitIndex, splitIndex + args.suffixLength); + + 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 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, + ); + if (shouldLogRecord) { + console.log("================================"); + console.log("PROMPT LENGTH:", promptLength); + + console.log("\nPROMPT IDS:"); + console.log(prompt.slice(0, 30)); + + console.log("\nGENERATED IDS:"); + console.log(generatedSuffix.slice(0, 30)); + + console.log("\nREFERENCE IDS:"); + console.log(reference.slice(0, 30)); + + console.log("\nPROMPT TEXT:"); + console.log(JSON.stringify(tokenizer.decode(prompt))); + + console.log("\nGENERATED TEXT:"); + console.log(JSON.stringify(tokenizer.decode(generatedSuffix))); + + console.log("\nREFERENCE TEXT:"); + console.log(JSON.stringify(tokenizer.decode(reference))); + + console.log("================================"); + } + + const exactMatch = + generatedSuffix.length === reference.length && + generatedSuffix.every((token, i) => token === reference[i]); + const bleu = bleu1to4(reference, generatedSuffix); + + results.push({ + recordIndex, + recordTokenLength: ids.length, + promptLength, + splitIndex, + exactMatch, + bleu, + memorizedByBleu: bleu > args.bleuThreshold, + promptText: tokenizer.decode(prompt), + referenceText: tokenizer.decode(reference), + generatedText: tokenizer.decode(generatedSuffix), + }); + + if (shouldLogRecord) { + console.log( + ` Done prompt length ${promptLength}: exact=${exactMatch}, BLEU=${bleu.toFixed(4)}`, + ); + } + } + + if (shouldLogRecord) { + console.log( + `Finished record ${recordIndex + 1}/${tokenizedRecords.length}; results so far: ${results.length}`, + ); + } + } + + 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, + decodingStrategy: args.decodingStrategy, + temperature: args.temperature, + topK: args.topK, + 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), + }; + + 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/aggregator/mean.spec.ts b/discojs/src/aggregator/mean.spec.ts index 6184171c6..3bfc367c4 100644 --- a/discojs/src/aggregator/mean.spec.ts +++ b/discojs/src/aggregator/mean.spec.ts @@ -1,3 +1,4 @@ +import * as tf from "@tensorflow/tfjs"; import { Set } from "immutable"; import { describe, expect, it } from "vitest"; import { WeightsContainer } from "#weights/index"; @@ -65,6 +66,29 @@ describe("mean aggregator", () => { expect(await WSIntoArrays(await results)).to.deep.equal([[1], [2]]); }); + it("aggregation leaves no dangling tensors", async () => { + const baseline = tf.memory().numTensors; + + const aggregator = new MeanAggregator(0, 2, "absolute"); + const [id1, id2] = ["client 1", "client 2"]; + aggregator.setNodes(Set.of(id1, id2)); + + const contribution1 = WeightsContainer.of([0], [1]); + const contribution2 = WeightsContainer.of([2], [3]); + + const result = aggregator.getPromiseForAggregation(); + aggregator.add(id1, contribution1, 0); + aggregator.add(id2, contribution2, 0); + expect(await WSIntoArrays(await result)).to.deep.equal([[1], [2]]); + + // the aggregator clones contributions on add and disposes the clones when + // aggregating; only the caller-owned inputs and the result should remain + contribution1.dispose(); + contribution2.dispose(); + (await result).dispose(); + expect(tf.memory().numTensors).to.equal(baseline); + }); + it("waits for 100% of the contributions by default", async () => { const aggregator = new MeanAggregator(); const [id1, id2] = ["client 1", "client 2"]; diff --git a/discojs/src/aggregator/mean.ts b/discojs/src/aggregator/mean.ts index 0f1faccfd..3d5854436 100644 --- a/discojs/src/aggregator/mean.ts +++ b/discojs/src/aggregator/mean.ts @@ -2,7 +2,7 @@ import type { Map } from "immutable"; import type { WeightsContainer } from "#weights/index"; import type { NodeID } from "#client/types"; -import { avg } from "#weights/index"; +import { avg as computeAvg } from "#weights/index"; import { AggregationStep } from "#aggregator/aggregator"; import type { ThresholdType } from "#aggregator/multiround"; @@ -10,6 +10,7 @@ import { MultiRoundAggregator } from "#aggregator/multiround"; /** * Mean aggregator whose aggregation step consists in computing the mean of the received weights. + * This aggregator extends MultiRoundAggregator while only performing a single round * */ export class MeanAggregator extends MultiRoundAggregator { @@ -23,13 +24,20 @@ export class MeanAggregator extends MultiRoundAggregator { } override _add(nodeId: 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 { @@ -39,8 +47,10 @@ export class MeanAggregator extends MultiRoundAggregator { this.log(AggregationStep.AGGREGATE); - const result = avg(currentContributions.values()); - return result; + const contributions = Array.from(currentContributions.values()); + const avg = computeAvg(contributions); + contributions.forEach((contribution) => contribution.dispose()); + return avg; } override makePayloads( diff --git a/discojs/src/aggregator/multiround.ts b/discojs/src/aggregator/multiround.ts index d435cfaad..43491a3cb 100644 --- a/discojs/src/aggregator/multiround.ts +++ b/discojs/src/aggregator/multiround.ts @@ -6,7 +6,7 @@ export type ThresholdType = "relative" | "absolute"; const debug = createDebug("discojs:aggregator:multiround"); /** - * Base class for multi-round aggregators. + * Base class for multi-round aggregators, with potentially only one round. * Multi-round aggregators are aggregators that wait for a certain number of contributions before aggregating. * They can be used to implement different aggregation strategies, such as Byzantine robust aggregation or Mean Aggregator. */ diff --git a/discojs/src/client/event_connection.ts b/discojs/src/client/event_connection.ts index 700b6994a..0a8aeed31 100644 --- a/discojs/src/client/event_connection.ts +++ b/discojs/src/client/event_connection.ts @@ -124,7 +124,10 @@ export class WebSocketServer validateReceived: (msg: unknown) => 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); @@ -144,8 +147,18 @@ export class WebSocketServer server.emit(msg.type, msg); }; + ws.onclose = (event) => { + 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 = () => { diff --git a/discojs/src/client/federated/federated_client.ts b/discojs/src/client/federated/federated_client.ts index 4f613a7cc..6d11f1081 100644 --- a/discojs/src/client/federated/federated_client.ts +++ b/discojs/src/client/federated/federated_client.ts @@ -90,7 +90,11 @@ export class FederatedClient extends Client<"federated"> { `[${shortenId(this.ownId)}] upon connecting, wait for participant flag %o`, this.waitingForMoreParticipants, ); - model.weights = weightsDecode(payload); + if (payload != null) { + const latestWeights = weightsDecode(payload); + model.weights = latestWeights; + } else debug(`[${shortenId(this.ownId)}] received an undefined payload`); + return model; } @@ -143,20 +147,34 @@ 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 msg: messages.SendPayload = { - type: MType.SendPayload, - payload: await weightsEncode(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 + const round = this.aggregator.round; + // block-scope the encoded payload so the potentially large buffer can be GC'd + // while we await the server's response below + { + const payload = await weightsEncode(payloadToServer); + debug( + "[%s] encoded payload for round %d byteLength=%d", + shortenId(this.ownId), + round, + payload.byteLength, + ); + + const msg: messages.SendPayload = { + type: MType.SendPayload, + payload, + round, + }; + + // Need to await the resulting global model right after sending our local contribution + // to make sure we don't miss it + this.server.send(msg); + } debug( - `[${shortenId(this.ownId)}] sent its local update to the server for round ${this.aggregator.round}`, + `[${shortenId(this.ownId)}] sent its local update to the server for round ${round}`, ); - this.server.send(msg); debug( - `[${shortenId(this.ownId)}] is waiting for server update for round ${this.aggregator.round + 1}`, + `[${shortenId(this.ownId)}] is waiting for server update for round ${round + 1}`, ); const { payload: payloadFromServer, diff --git a/discojs/src/client/federated/messages.ts b/discojs/src/client/federated/messages.ts index d45e7c542..9172efb56 100644 --- a/discojs/src/client/federated/messages.ts +++ b/discojs/src/client/federated/messages.ts @@ -21,7 +21,7 @@ export interface NewFederatedNodeInfo { type: MType.NewFederatedNodeInfo; id: NodeID; waitForMoreParticipants: boolean; - payload: serialization.Encoded; + payload?: serialization.Encoded | null; round: number; nbOfParticipants: number; } diff --git a/discojs/src/client/local_client.ts b/discojs/src/client/local_client.ts index 24e4eb59e..06ca66886 100644 --- a/discojs/src/client/local_client.ts +++ b/discojs/src/client/local_client.ts @@ -9,10 +9,11 @@ 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())); } } diff --git a/discojs/src/default_tasks/goldfish.ts b/discojs/src/default_tasks/goldfish.ts new file mode 100644 index 000000000..a23390d63 --- /dev/null +++ b/discojs/src/default_tasks/goldfish.ts @@ -0,0 +1,49 @@ +import type { TaskProvider } from "#task/index"; +import { Tokenizer, cards } from "#models/index"; + +export const goldfish: TaskProvider<"text", "federated"> = { + async getTask() { + return { + id: "goldfish", + dataType: "text", + displayInformation: { + title: "Privacy-Preserving Fine-tuning of GPT-2", + summary: { + preview: + "Fine-tune a pre-trained GPT model collaboratively and privately with the Goldfish loss.", + overview: + "Fine-tune a pre-trained GPT-2 model created by the ONNX converter in your browser collaboratively without sharing your raw data. The model is loaded from Google Cloud Storage and fine-tuned using federated learning.", + }, + model: [ + "The model is a pre-trained GPT-2 architecture converted from ONNX and loaded from Google Cloud Storage.", + "The tokenizer used for preprocessing is the GPT-2 Byte-Pair encoding tokenizer.", + "The model is trained via an Adam optimizer with unit gradient clipping and softmax cross-entropy loss.", + "Context length is kept at 512 to match the pre-trained model, with batch size at 8.", + ].join(" "), + dataFormatInformation: + "You can use any natural language (text) dataset. The dataset should be formatted as a plain text file with each line representing a segment of text.", + dataExample: + "For the first twenty years of its existence , the only staged performances of Parsifal took place in the Bayreuth Festspielhaus , the venue for which Wagner conceived the work.", + }, + trainingInformation: { + scheme: "federated", + aggregationStrategy: "mean", + minNbOfParticipants: 2, + 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", + goldfishLoss: { + enabled: true, + k: 4, + h: 13, + }, + }, + }; + }, + modelCard: cards.Goldfish, +}; diff --git a/discojs/src/default_tasks/index.ts b/discojs/src/default_tasks/index.ts index 6c2714615..036a11046 100644 --- a/discojs/src/default_tasks/index.ts +++ b/discojs/src/default_tasks/index.ts @@ -4,3 +4,4 @@ export { mnist } from "./mnist.js"; export { titanic } from "./titanic.js"; export { wikitext } from "./wikitext.js"; export { tinderDog } from "./tinder_dog.js"; +export { goldfish } from "./goldfish.js"; diff --git a/discojs/src/index.ts b/discojs/src/index.ts index d9392ea20..9d3612710 100644 --- a/discojs/src/index.ts +++ b/discojs/src/index.ts @@ -53,7 +53,11 @@ export { evaluate_hellaswag, DefaultGenerationConfig, } from "./models/index.js"; -export type { GPTConfig, HellaSwagExample } from "./models/index.js"; +export type { + GPTConfig, + HellaSwagExample, + GoldfishLossConfig, +} from "./models/index.js"; export { EventEmitter } from "./utils/event_emitter.js"; diff --git a/discojs/src/models/cards/goldfishClassifier.ts b/discojs/src/models/cards/goldfishClassifier.ts new file mode 100644 index 000000000..83ba46a88 --- /dev/null +++ b/discojs/src/models/cards/goldfishClassifier.ts @@ -0,0 +1,49 @@ +import type { Model } from "#models/model"; +import type { ModelCard } from "#models/model_card"; +import { GPT } from "#models/implementations/index"; +import { decode as modelDecode } from "#serialization/model"; + +export const Goldfish: ModelCard<"text"> = { + card: { + id: "goldfish", + name: "Privacy-Preserving Fine-tuning of GPT-2", + dataType: "text", + contextLength: 512, + }, + + async getModel(): Promise> { + // 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"; + + 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 modelDecode(encodedData); + + if (!(model instanceof 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/models/cards/index.ts b/discojs/src/models/cards/index.ts index acff0bf71..98b0b3585 100644 --- a/discojs/src/models/cards/index.ts +++ b/discojs/src/models/cards/index.ts @@ -3,4 +3,5 @@ export { CIFAR10Classifier } from "./CIFAR10Classifier.js"; export { MNISTClassifier } from "./MNISTClassifier.js"; export { DogClassifier } from "./dogClassifier.js"; export { TitanicClassifier } from "./titanicClassifier.js"; -export { Wikitext } from "./wikitext.js"; +export { Wikitext } from "./wikitextClassifier.js"; +export { Goldfish } from "./goldfishClassifier.js"; diff --git a/discojs/src/models/cards/wikitext.ts b/discojs/src/models/cards/wikitextClassifier.ts similarity index 100% rename from discojs/src/models/cards/wikitext.ts rename to discojs/src/models/cards/wikitextClassifier.ts diff --git a/discojs/src/models/implementations/gpt/config.ts b/discojs/src/models/implementations/gpt/config.ts index 247bfec15..fb178e987 100644 --- a/discojs/src/models/implementations/gpt/config.ts +++ b/discojs/src/models/implementations/gpt/config.ts @@ -27,11 +27,19 @@ 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, 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, diff --git a/discojs/src/models/implementations/gpt/gpt.ts b/discojs/src/models/implementations/gpt/gpt.ts index 186279178..280895038 100644 --- a/discojs/src/models/implementations/gpt/gpt.ts +++ b/discojs/src/models/implementations/gpt/gpt.ts @@ -16,7 +16,10 @@ import { Model } from "#models/model"; import { GPTModel } from "#models/implementations/gpt/model"; import evaluate from "#models/implementations/gpt/evaluate"; import { DefaultGPTConfig } from "#models/implementations/gpt/config"; -import type { GPTConfig } from "#models/implementations/gpt/config"; +import type { + GPTConfig, + GoldfishLossConfig, +} from "#models/implementations/gpt/config"; import { DefaultGenerationConfig } from "#models/generation"; import type { GenerationConfig } from "#models/generation"; @@ -34,6 +37,7 @@ export class GPT extends Model<"text"> { readonly #contextLength: number; readonly #maxBatchCount: number; readonly #vocabSize: number; + #iterationCount = 0; constructor( partialConfig?: Partial, @@ -51,6 +55,18 @@ export class GPT extends Model<"text"> { this.#vocabSize = partialConfig?.vocabSize ?? DefaultGPTConfig.vocabSize; } + setDebugLabel(label: string): void { + this.model.setDebugLabel(label); + } + + setGoldfishLoss(config: GoldfishLossConfig | undefined): void { + 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. @@ -70,7 +86,40 @@ 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); + } + + const validation = + validationDataset && (await this.evaluate(validationDataset)); + epochTime = performance.now() - epochTime; + + 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); @@ -79,18 +128,21 @@ export class GPT extends Model<"text"> { 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/implementations/gpt/model.ts b/discojs/src/models/implementations/gpt/model.ts index ebee4e92c..8ff763200 100644 --- a/discojs/src/models/implementations/gpt/model.ts +++ b/discojs/src/models/implementations/gpt/model.ts @@ -1,7 +1,10 @@ import createDebug from "debug"; import * as tf from "@tensorflow/tfjs"; -import type { GPTConfig } from "#models/implementations/gpt/config"; +import type { + GPTConfig, + GoldfishLossConfig, +} from "#models/implementations/gpt/config"; import { getModelSizes, DefaultGPTConfig, @@ -15,6 +18,18 @@ import { GPTArchitecture } from "#models/implementations/gpt/layers"; 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... */ @@ -33,6 +48,8 @@ export declare abstract class Dataset { */ export class GPTModel extends tf.LayersModel { protected readonly config: Required; + #debugLabel?: string; + #goldfishLoss?: GoldfishLossConfig; constructor( partialConfig?: Partial, @@ -67,6 +84,20 @@ export class GPTModel extends tf.LayersModel { return this.config; } + setDebugLabel(label: string): void { + 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}`; + } + override compile() { if (this.optimizer !== undefined) return; this.optimizer = @@ -75,25 +106,40 @@ 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, + 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; + + 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 }; @@ -103,25 +149,29 @@ export class GPTModel extends tf.LayersModel { 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 goldfishLoss = this.#goldfishLoss; + const goldfishMask = + goldfishLoss === undefined + ? undefined + : this.#buildGoldfishMask(xs, goldfishLoss); const lossTensor = tf.tidy(() => { const { grads, value: lossTensor } = this.optimizer.computeGradients( @@ -131,39 +181,50 @@ 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(); + lossTensor.dispose(); + tf.dispose([xs, ys]); averageLoss += loss; weightUpdateTime = performance.now() - weightUpdateTime; - tf.dispose([xs, ys, lossTensor]); - 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); + 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, loss, memory, allocated: tf.memory().numTensors, + processMemory: processMemory(), preprocessingTime, weightUpdateTime, }); @@ -190,4 +251,74 @@ 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), + ); + + 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"); + } + + /** + * Computes a deterministic 32-bit FNV-1a-style hash of a token context. + * Each token ID is mixed in little-endian byte order, followed by a separator + * byte to preserve token boundaries. Goldfish uses the unsigned result modulo + * `k` to consistently select which target-token losses to drop. + */ + #hashTokenContext(tokens: number[]): number { + let hash = 0x811c9dc5; + + 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/implementations/index.ts b/discojs/src/models/implementations/index.ts index 982195083..03e857431 100644 --- a/discojs/src/models/implementations/index.ts +++ b/discojs/src/models/implementations/index.ts @@ -1,4 +1,4 @@ export { GPT } from "./gpt/index.js"; -export type { GPTConfig } from "./gpt/config.js"; +export type { GPTConfig, GoldfishLossConfig } from "./gpt/config.js"; export type { HellaSwagDataset, HellaSwagExample } from "./hellaswag.js"; export { evaluate as evaluate_hellaswag, HELLASWAG_URL } from "./hellaswag.js"; diff --git a/discojs/src/models/index.ts b/discojs/src/models/index.ts index a7a5cb0da..e49ca23b3 100644 --- a/discojs/src/models/index.ts +++ b/discojs/src/models/index.ts @@ -7,6 +7,7 @@ export type { GenerationConfig } from "./generation.js"; export type { GPTConfig, + GoldfishLossConfig, HellaSwagDataset, HellaSwagExample, } from "./implementations/index.js"; diff --git a/discojs/src/privacy.spec.ts b/discojs/src/privacy.spec.ts index 9545db314..df941c82b 100644 --- a/discojs/src/privacy.spec.ts +++ b/discojs/src/privacy.spec.ts @@ -23,6 +23,17 @@ describe("frobeniusNorm", () => { const n = await frobeniusNorm(t); expect(n).toBeCloseTo(5, 1e-12); }); + + it("does not leak intermediate tensors nor dispose its input", async () => { + const baseline = tf.memory().numTensors; + + const t = tf.tensor([3, 4]); + await frobeniusNorm(t); + + expect(t.isDisposed).toBe(false); + t.dispose(); + expect(tf.memory().numTensors).toBe(baseline); + }); }); describe("clipNorm", () => { @@ -46,6 +57,22 @@ describe("clipNorm", () => { [0, 3], ]); }); + + it("does not leak intermediate tensors nor dispose its input", async () => { + const baseline = tf.memory().numTensors; + + // one layer above the radius (clipped), one within it (kept as-is) + const input = WeightsContainer.of([2], [0, 6]); + const result = await clipNorm(input, [1, 10]); + + for (const weight of input.weights) expect(weight.isDisposed).toBe(false); + + // the result must own fresh tensors, not views of the input + input.dispose(); + expect(await WSIntoArrays(result)).toEqual([[1], [0, 6]]); + result.dispose(); + expect(tf.memory().numTensors).toBe(baseline); + }); }); describe("addOptimalNoise", () => { @@ -69,6 +96,24 @@ describe("addOptimalNoise", () => { expect(Number.isFinite(resultArrays[1][0])).toBe(true); expect(Number.isFinite(resultArrays[1][1])).toBe(true); }); + + it("does not leak intermediate tensors nor dispose its input", async () => { + const baseline = tf.memory().numTensors; + + const input = WeightsContainer.of([3, 4], [0, 6]); + const result = await addOptimalNoise(input, 1, 1e-5, [5, 3]); + + for (const weight of input.weights) expect(weight.isDisposed).toBe(false); + + // the internally clipped weights must be disposed and the result must own + // fresh tensors, not views of the input + input.dispose(); + expect((await WSIntoArrays(result)).flat().every(Number.isFinite)).toBe( + true, + ); + result.dispose(); + expect(tf.memory().numTensors).toBe(baseline); + }); }); describe("getClippingRadius", () => { diff --git a/discojs/src/privacy.ts b/discojs/src/privacy.ts index 9829b694c..3ddd76b0d 100644 --- a/discojs/src/privacy.ts +++ b/discojs/src/privacy.ts @@ -6,7 +6,9 @@ import type { WeightNormHistory } from "#training/types"; /** 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]); } @@ -52,9 +54,13 @@ export async function addOptimalNoise( ); 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/processing/processing.spec.ts b/discojs/src/processing/processing.spec.ts index bf995b99b..a20e366a5 100644 --- a/discojs/src/processing/processing.spec.ts +++ b/discojs/src/processing/processing.spec.ts @@ -5,6 +5,12 @@ import { preprocess } from "#processing/processing"; import type { Task } from "#task/index"; import { Dataset } from "#dataset/index"; +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/processing.ts b/discojs/src/processing/processing.ts index 3cc9e981c..bf7c7fb7a 100644 --- a/discojs/src/processing/processing.ts +++ b/discojs/src/processing/processing.ts @@ -51,6 +51,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] >; diff --git a/discojs/src/serialization/model.ts b/discojs/src/serialization/model.ts index 8a56e524c..c1f7f2bef 100644 --- a/discojs/src/serialization/model.ts +++ b/discojs/src/serialization/model.ts @@ -1,8 +1,10 @@ import type tf from "@tensorflow/tfjs"; import { encode as w_encode, decode as w_decode } from "#serialization/weights"; -import { GPT, TFJS } from "#models/index"; -import type { Model, GPTConfig } from "#models/index"; +import { GPT } from "#models/implementations/index"; +import type { GPTConfig } from "#models/implementations/index"; +import { TFJS } from "#models/tfjs"; +import type { Model } from "#models/model"; import type { DataType } from "#types/index"; import type { Encoded } from "#serialization/coder"; @@ -41,12 +43,14 @@ export async function decode(encoded: Encoded): Promise> { "invalid encoding, encoding isn't an array or doesn't contain enough values", ); } + const type = raw[0] as unknown; if (typeof type !== "number") { throw new Error( "invalid encoding, first encoding field should be the model type", ); } + const rawModel = raw[1] as unknown; switch (type) { case Type.TFJS: { diff --git a/discojs/src/task/training_information.ts b/discojs/src/task/training_information.ts index be7e4bd01..4f53ec62e 100644 --- a/discojs/src/task/training_information.ts +++ b/discojs/src/task/training_information.ts @@ -66,6 +66,13 @@ 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. + // roundDuration is ignored if roundIterations is set + roundIterations: z.number().positive().int().optional(), + // run validation every N aggregation rounds. If 0, validation metrics are skipped. + validationFrequency: z.number().nonnegative().int().optional(), + // whether to validate before aggregation, after aggregation, or at both points + validationMode: z.enum(["before", "after", "both"]).optional(), // fraction of data to keep for validation, note this only works for image data validationSplit: z.number().min(0).max(1), // batch size of training data @@ -96,6 +103,16 @@ 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(), + learningRate: z.number().positive().optional(), }), } satisfies Record; diff --git a/discojs/src/training/disco.ts b/discojs/src/training/disco.ts index 0abd4a38e..96e0f76ca 100644 --- a/discojs/src/training/disco.ts +++ b/discojs/src/training/disco.ts @@ -15,14 +15,17 @@ import { EventEmitter } from "#utils/event_emitter"; import * as clients from "#client/index"; import * as processing from "#processing/index"; import * as async_iterator from "#utils/async_iterator"; - +import type { GoldfishLossConfig } from "#models/implementations/gpt/config"; import type { RoundLogs } from "#training/trainer"; import { Trainer } from "#training/trainer"; import type { RoundStatus, SummaryLogs } from "#training/types"; +import createDebug from "debug"; +const debug = createDebug("discojs:training:disco"); interface DiscoConfig { scheme: N; logger: Logger; + debugLabel?: string; /** * keep preprocessed dataset in memory while training @@ -51,6 +54,9 @@ function buildSummaryLog( roundValidationAccuracy: roundLogs.preRoundValidation?.accuracy, validationLoss: epochLogs.validation?.loss, validationAccuracy: epochLogs.validation?.accuracy, + postAggregationValidationLoss: roundLogs.postAggregationValidation?.loss, + postAggregationValidationAccuracy: + roundLogs.postAggregationValidation?.accuracy, }; } @@ -68,6 +74,8 @@ export class Disco extends EventEmitter<{ readonly #logger: Logger; readonly #task: Task; readonly #preprocessOnce: boolean; + // Forwarded to compatible models to identify this client in debug output. + readonly #debugLabel?: string; /** * Connect to the given task and get ready to train. @@ -85,7 +93,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(), @@ -111,6 +119,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); @@ -148,21 +157,27 @@ 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 }> = []; + 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); @@ -180,8 +195,11 @@ 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); } @@ -193,19 +211,27 @@ 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 + debug("Connecting to client and fetching initial model..."); + // TODO unsafe cast this.trainer.model = (await this.#client.connect()) as Model; + this.#setModelDebugLabel(this.trainer.model); + this.#setModelTrainingOptions(this.trainer.model); + debug("Initial model fetched successfully"); 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); @@ -239,15 +265,27 @@ 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"), ); } + this.#logger.success( + [ + `Round: ${roundNum}`, + roundLogs.postAggregationValidation !== undefined + ? `Post-aggregation loss: ${roundLogs.postAggregationValidation.loss}` + : "", + roundLogs.postAggregationValidation + ? `Post-aggregation accuracy: ${roundLogs.postAggregationValidation.accuracy}` + : "", + ].join("\n"), + ); + return roundLogs; }.bind(this)(); } @@ -261,6 +299,44 @@ 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); + } + + #setModelTrainingOptions(model: Model): void { + if (this.#task.dataType !== "text") return; + + const configurableModel = model as Model & { + setGoldfishLoss?: (config: GoldfishLossConfig | undefined) => void; + setLearningRate?: (learningRate: number) => void; + }; + + configurableModel.setGoldfishLoss?.( + this.#task.trainingInformation.goldfishLoss, + ); + if (this.#task.trainingInformation.goldfishLoss?.enabled === true) { + const { k, h, padTokenId } = this.#task.trainingInformation.goldfishLoss; + debug( + `Using Goldfish loss with k=${k}, h=${h}` + + (padTokenId === undefined ? "" : `, padTokenId=${padTokenId}`), + ); + } + if (this.#task.trainingInformation.learningRate !== undefined) { + configurableModel.setLearningRate?.( + this.#task.trainingInformation.learningRate, + ); + debug( + `Using GPT learning rate ${this.#task.trainingInformation.learningRate}`, + ); + } + } + async #preprocessSplitAndBatch( dataset: Dataset, ): Promise< @@ -286,6 +362,41 @@ 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) diff --git a/discojs/src/training/trainer.ts b/discojs/src/training/trainer.ts index 023b03766..4a18bf90e 100644 --- a/discojs/src/training/trainer.ts +++ b/discojs/src/training/trainer.ts @@ -3,10 +3,9 @@ import { List, Repeat } from "immutable"; import type { Model } from "#models/index"; import type { DataFormat, DataType, Network } from "#types/index"; -import type { Batched } from "#dataset/index"; +import type { Batched, Dataset } from "#dataset/index"; import type { Task } from "#task/index"; -import type { Dataset } from "#dataset/index"; import type { BatchLogs, EpochLogs, ValidationMetrics } from "#models/index"; import { WeightsContainer } from "#weights/index"; import type { Client } from "#client/index"; @@ -14,12 +13,18 @@ import type { Client } from "#client/index"; import * as async_iterator from "#utils/async_iterator"; import * as privacy from "#root/privacy"; -import type { WeightNormHistory } from "#training/types"; +import type { + WeightNormHistory, + IterationTrainableTextModel, +} from "#training/types"; +import createDebug from "debug"; +const debug = createDebug("discojs:training:trainer"); export interface RoundLogs { epochs: List; participants: number; preRoundValidation?: ValidationMetrics; + postAggregationValidation?: ValidationMetrics; } /** List of weight update norms */ @@ -49,6 +54,9 @@ export class Trainer { AsyncGenerator, RoundLogs>, void >; + readonly #roundIterations?: number; + readonly #validationFrequency?: number; + readonly #validationMode: "before" | "after" | "both"; // Map of weight Index and weight update #weightNormHistory: WeightNormHistory = List(); #previousRoundWeights?: WeightsContainer; @@ -67,10 +75,36 @@ export class Trainer { this.#client = client; this.#roundDuration = task.trainingInformation.roundDuration; this.#epochs = task.trainingInformation.epochs; + this.#roundIterations = task.trainingInformation.roundIterations; + this.#validationFrequency = task.trainingInformation.validationFrequency; + this.#validationMode = task.trainingInformation.validationMode ?? "before"; 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 ( + this.#validationFrequency !== undefined && + (!Number.isInteger(this.#validationFrequency) || + this.#validationFrequency < 0) + ) + throw new Error("validationFrequency must be a non-negative integer"); + + 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}`, ); @@ -87,20 +121,28 @@ 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", ); try { - this.#training = this.#runRounds(dataset, validationDataset); + this.#training = + this.#roundIterations === undefined + ? this.#runRoundsByEpoch(dataset, validationDataset) + : this.#runRoundsByIteration(dataset, validationDataset); yield* this.#training; } finally { this.#training = undefined; } } - async *#runRounds( + /** + * Runs epoch-based training, aggregating after `roundDuration` complete + * passes over the training dataset until the configured epochs are reached. + */ + async *#runRoundsByEpoch( dataset: Dataset>, validationDataset?: Dataset>, ): AsyncGenerator< @@ -108,6 +150,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(); @@ -116,65 +161,269 @@ export class Trainer { this.model.weights.weights.map((t) => t.clone()), ); - yield this.#runRound(dataset, validationDataset); + const roundValidationDataset = this.#shouldValidateRound(round) + ? validationDataset + : undefined; - let roundWeights = this.model.weights; + yield this.#runRoundByEpoch( + dataset, + this.#shouldValidateBeforeAggregation() + ? roundValidationDataset + : undefined, + this.#shouldValidateAfterAggregation() + ? roundValidationDataset + : undefined, + totalRound, + ); + } + } - // 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, + /** + * Runs iteration-based training, aggregating after `roundIterations` + * batches while preserving the dataset iterator between rounds. A new + * iterator is created only when the next configured epoch begins. + */ + async *#runRoundsByIteration( + dataset: Dataset>, + validationDataset?: Dataset>, + ): AsyncGenerator< + AsyncGenerator, RoundLogs>, + void + > { + 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](); + let next = await trainingIterator.next(); + while (next.done !== true) { + await this.#client.onRoundBeginCommunication(); + + this.#previousRoundWeights = new WeightsContainer( + this.model.weights.weights.map((t) => t.clone()), ); - roundWeights = await applyOptimalPrivacy( - this.#previousRoundWeights, - roundWeights, - this.#privacy, - this.#weightNormHistory, + let firstBatch: Batched | undefined = + next.value; + let done = false; + const prefixedIterator: AsyncIterator< + Batched + > = { + next: async () => { + if (firstBatch !== undefined) { + const value = firstBatch; + firstBatch = undefined; + return { value, done: false }; + } + + return await trainingIterator.next(); + }, + }; + + const roundValidationDataset = this.#shouldValidateRound(round) + ? validationDataset + : undefined; + + yield this.#runRoundByIteration( + prefixedIterator, + this.#roundIterations, + this.#shouldValidateBeforeAggregation() + ? roundValidationDataset + : undefined, + this.#shouldValidateAfterAggregation() + ? roundValidationDataset + : undefined, totalRound, + (roundDone) => (done = roundDone), ); - } - // Get the updated weights - const networkWeights = - await this.#client.onRoundEndCommunication(roundWeights); - // Update the local weights - this.model.weights = networkWeights; + round++; + if (done) break; + next = await trainingIterator.next(); + } } } - async *#runRound( + /** + * Trains one epoch-based round by making `roundDuration` complete passes + * over the dataset, then exchanges weights and returns the round metrics. + */ + async *#runRoundByEpoch( dataset: Dataset>, - validationDataset?: Dataset>, + preAggregationValidationDataset: + | Dataset> + | undefined, + postAggregationValidationDataset: + | Dataset> + | undefined, + totalRound: number, ): AsyncGenerator, RoundLogs> { let epochsLogs = List(); + debug("Run round"); + // Before starting the training, get the validation of global model const validation = - validationDataset !== undefined - ? await this.model.evaluate(validationDataset) + preAggregationValidationDataset !== undefined + ? await this.model.evaluate(preAggregationValidationDataset) : undefined; for (let epoch = 0; epoch < this.#roundDuration; epoch++) { const [gen, epochLogs] = async_iterator.split( - this.model.train(dataset, validationDataset), + this.model.train(dataset, preAggregationValidationDataset), ); yield gen; epochsLogs = epochsLogs.push(await epochLogs); } + const participants = this.#client.nbOfParticipants; + const postAggregationValidation = await this.#finishRoundCommunication( + totalRound, + postAggregationValidationDataset, + ); + return { epochs: epochsLogs, - participants: this.#client.nbOfParticipants, + participants, + preRoundValidation: validation, + postAggregationValidation, + }; + } + + /** + * Trains one iteration-based round by consuming at most `maxBatchCount` + * batches from the supplied iterator without rewinding it, then exchanges + * weights and returns the round metrics. + */ + async *#runRoundByIteration( + datasetIterator: AsyncIterator>, + maxBatchCount: number, + preAggregationValidationDataset: + | Dataset> + | undefined, + postAggregationValidationDataset: + | Dataset> + | undefined, + totalRound: number, + setDone?: (done: boolean) => void, + ): AsyncGenerator, RoundLogs> { + const model = this.model as unknown as IterationTrainableTextModel; + if (typeof model.trainNextBatches !== "function") + throw new Error("model does not support iteration-based training"); + + debug("Run iteration-based round"); + let iterationLogs = List(); + + const validation = + preAggregationValidationDataset !== undefined + ? await this.model.evaluate(preAggregationValidationDataset) + : undefined; + + const [gen, epochLogs] = async_iterator.split( + model.trainNextBatches( + datasetIterator as AsyncIterator< + Batched + >, + maxBatchCount, + preAggregationValidationDataset as + | Dataset> + | undefined, + setDone, + ), + ); + + yield gen; + iterationLogs = iterationLogs.push(await epochLogs); + + const participants = this.#client.nbOfParticipants; + const postAggregationValidation = await this.#finishRoundCommunication( + totalRound, + postAggregationValidationDataset, + ); + + return { + epochs: iterationLogs, + participants, preRoundValidation: validation, + postAggregationValidation, }; } + + #shouldValidateBeforeAggregation(): boolean { + return this.#validationMode !== "after"; + } + + #shouldValidateAfterAggregation(): boolean { + return this.#validationMode !== "before"; + } + + #shouldValidateRound(round: number): boolean { + if (this.#validationFrequency === undefined) return true; + if (this.#validationFrequency === 0) return false; + return round % this.#validationFrequency === 0; + } + + async #finishRoundCommunication( + totalRound: number, + validationDataset?: Dataset>, + ): Promise { + let roundWeights = this.model.weights; + let disposeRoundWeightsAfterSend = false; + + 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); + 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 = privateRoundWeights; + disposeRoundWeightsAfterSend = true; + } + + const networkWeights = + await this.#client.onRoundEndCommunication(roundWeights); + this.model.weights = networkWeights; + + return validationDataset !== undefined + ? await this.model.evaluate(validationDataset) + : undefined; + } finally { + if (disposeRoundWeightsAfterSend) roundWeights.dispose(); + this.#previousRoundWeights?.dispose(); + this.#previousRoundWeights = undefined; + } + } } /** ALDP-FL implementation */ @@ -192,25 +441,9 @@ async function applyOptimalPrivacy( totalRound: number, ): Promise { let ret = current; - - // Clipping radius for BFT - if ("byzantineFaultTolerance" in options) { - // might need to change the variable name - const previousRoundWeights = - previous ?? current.map((w) => tf.zerosLike(w)); - const weightsProgress = current.sub(previousRoundWeights); - ret = previousRoundWeights.add( - await privacy.clipNorm( - weightsProgress, - Repeat(options.byzantineFaultTolerance.clippingRadius) - .take(weightsProgress.weights.length) - .toArray(), - ), - ); - } + const dpOptions = options.differentialPrivacy; // Adding Gaussian noise for DP - const dpOptions = options.differentialPrivacy; if (dpOptions !== undefined) { const dpDefaultRadius = dpOptions.clippingRadius; // options.dpDefaultClippingRadius should be a number @@ -235,14 +468,52 @@ async function applyOptimalPrivacy( ) : dpClippingRadius; - ret = previousEpochWeights.add( - await privacy.addOptimalNoise( - weightsProgress, - epsilon, - delta, - effectiveRadius, - ), + 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(); + } + } + // Clipping radius for BFT if DP didn't already clip + else if ("byzantineFaultTolerance" in options) { + // might need to change the variable name + const previousRoundWeights = + previous ?? current.map((w) => tf.zerosLike(w)); + const weightsProgress = current.sub(previousRoundWeights); + const clippedProgress = await privacy.clipNorm( + weightsProgress, + Repeat(options.byzantineFaultTolerance.clippingRadius) + .take(weightsProgress.weights.length) + .toArray(), + ); + try { + ret = previousRoundWeights.add(clippedProgress); + } finally { + weightsProgress.dispose(); + clippedProgress.dispose(); + if (previous === undefined) previousRoundWeights.dispose(); + } } return ret; } diff --git a/discojs/src/training/types.ts b/discojs/src/training/types.ts index 52459d021..d39605c4a 100644 --- a/discojs/src/training/types.ts +++ b/discojs/src/training/types.ts @@ -1,7 +1,20 @@ import type { List } from "immutable"; +import type { DataFormat } from "#types/index"; +import type { Batched, Dataset } from "#dataset/index"; +import type { Model } from "#models/index"; +import type { BatchLogs, EpochLogs } from "#models/index"; export type WeightNormHistory = List>; +export type IterationTrainableTextModel = Model<"text"> & { + trainNextBatches( + trainingIterator: AsyncIterator>, + maxBatchCount: number, + validationDataset?: Dataset>, + setDone?: (done: boolean) => void, + ): AsyncGenerator; +}; + export type SummaryLogs = { round: number; epoch: number; @@ -13,6 +26,8 @@ export type SummaryLogs = { roundValidationAccuracy?: number; validationLoss?: number; validationAccuracy?: number; + postAggregationValidationLoss?: number; + postAggregationValidationAccuracy?: number; }; export type RoundStatus = diff --git a/discojs/src/weights/aggregation.spec.ts b/discojs/src/weights/aggregation.spec.ts index 78d2e39ab..cf86e622a 100644 --- a/discojs/src/weights/aggregation.spec.ts +++ b/discojs/src/weights/aggregation.spec.ts @@ -1,3 +1,4 @@ +import * as tf from "@tensorflow/tfjs"; import { assert, describe, it } from "vitest"; import { WeightsContainer } from "#weights/weights_container"; import { avg, sum, diff } from "#weights/aggregation"; @@ -14,6 +15,56 @@ describe("weights aggregation", () => { assert.isTrue(actual.equals(expected)); }); + it("avg does not leak intermediate tensors", () => { + const baseline = tf.memory().numTensors; + + const inputs = [ + WeightsContainer.of([1, 2, 3, -1], [-5, 6]), + WeightsContainer.of([2, 3, 7, 1], [-10, 5]), + WeightsContainer.of([3, 1, 5, 3], [-15, 19]), + ]; + const result = avg(inputs); + + inputs.forEach((input) => input.dispose()); + result.dispose(); + assert.strictEqual(tf.memory().numTensors, baseline); + }); + + it("avg of a single container does not leak intermediate tensors", async () => { + const baseline = tf.memory().numTensors; + + const input = WeightsContainer.of([1, 2], [3]); + const result = avg([input]); + + // read values without allocating comparison tensors + assert.deepStrictEqual( + await Promise.all(result.weights.map(async (w) => [...(await w.data())])), + [[1, 2], [3]], + ); + input.dispose(); + result.dispose(); + assert.strictEqual(tf.memory().numTensors, baseline); + }); + + it("avg does not dispose nor alias its inputs", async () => { + const inputs = [ + WeightsContainer.of([1, 2], [3]), + WeightsContainer.of([3, 4], [5]), + ]; + const result = avg(inputs); + + for (const input of inputs) + for (const weight of input.weights) assert.isFalse(weight.isDisposed); + + // the result must own fresh tensors, not views of the inputs + inputs.forEach((input) => input.dispose()); + assert.deepStrictEqual( + await Promise.all(result.weights.map(async (w) => [...(await w.data())])), + [[2, 3], [4]], + ); + result.dispose(); + }); + it("sum of weights with two operands", () => { const actual = sum([ [[3, -4], [9]], diff --git a/discojs/src/weights/aggregation.ts b/discojs/src/weights/aggregation.ts index 58773d13e..995d54dc5 100644 --- a/discojs/src/weights/aggregation.ts +++ b/discojs/src/weights/aggregation.ts @@ -67,6 +67,20 @@ export function diff( export function avg( weights: Iterable, ): WeightsContainer { - const ws = List(weights); - return sum(ws).map((w) => w.div(ws.size)); + const ws = parseWeights(weights); + const first = ws.first(); + if (first === undefined) throw new Error("no weights to work with"); + let summed: WeightsContainer = first.map((weight) => weight.clone()); + + try { + for (const weights of ws.rest()) { + const next: WeightsContainer = summed.add(weights); + summed.dispose(); + summed = next; + } + + return summed.map((weight) => weight.div(ws.size)); + } finally { + summed.dispose(); + } } diff --git a/onnx-converter/src/convert_onnx.ts b/onnx-converter/src/convert_onnx.ts index 505860bac..39a6b98e8 100644 --- a/onnx-converter/src/convert_onnx.ts +++ b/onnx-converter/src/convert_onnx.ts @@ -7,6 +7,7 @@ import { GPT, modelEncode } 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"; @@ -32,8 +33,10 @@ async function main() { console.log("ONNX model loaded successfully"); // Init empty TF.js model - // Context length value from https://huggingface.co/Xenova/gpt2/blob/main/config.json - const gptModel = new GPT({ modelType: "gpt2", contextLength: 1024 }); + const gptModel = new 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,7 +62,18 @@ 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); } diff --git a/package.json b/package.json index 8e13c87ea..1817555bb 100644 --- a/package.json +++ b/package.json @@ -5,7 +5,7 @@ "lint": "eslint", "format:check": "prettier -c .", "format:fix": "prettier -w --list-different .", - "check_cycles": "dpdm --tsconfig discojs/tsconfig.lib.json --circular --exit-code circular:1 'discojs/src/**/*.ts'" + "check_cycles": "dpdm --tsconfig discojs/tsconfig.lib.json --circular --no-warning --no-tree --exit-code circular:1 'discojs/src/**/*.ts'" }, "dependencies": { "@tensorflow/tfjs-node": "catalog:" diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 0937a0e65..1a9b0dab7 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -7,52 +7,52 @@ importers: configDependencies: {} packageManagerDependencies: '@pnpm/exe': - specifier: '>=11.0.0 <12.0.0' - version: 11.9.0 + specifier: '>=11.0.0 <11.12.0' + version: 11.11.0 pnpm: - specifier: '>=11.0.0 <12.0.0' - version: 11.9.0 + specifier: '>=11.0.0 <11.12.0' + version: 11.11.0 packages: - '@pnpm/exe@11.9.0': - resolution: {integrity: sha512-pPPOpR79qW3nsNhlyDIdfstli4Bi78mk8r22ySxpFRwMbO8KXSjGrVzGmJBsVX39NnJTh7/WADj527nZhG9H9g==} + '@pnpm/exe@11.11.0': + resolution: {integrity: sha512-OpTrbkAU0Ur8gBwhAT6mbWwlA1bFU8zPcFNdp/img/RqFIc7Dn0j0crW9rz5hxTWo/UBQjy6Sb9jh+N5bRSd+g==} hasBin: true - '@pnpm/linux-arm64@11.9.0': - resolution: {integrity: sha512-XYmY2qadHauBA3QaHi2R7fI6kt5Flje0WHz9MVrbH0kVH/XLpfOLnwPeE1+EX6K/nDa2CBvzp35VjYCNGFJa9A==} + '@pnpm/linux-arm64@11.11.0': + resolution: {integrity: sha512-62K/kQY3jaoQ96MNUi0NLcTSSDO2QrtziOA2IK5R/m+DxpSwXbxFSQHYx8oOg0r1+MFVqwuDkFjXe6DyW3y58g==} cpu: [arm64] os: [linux] - '@pnpm/linux-x64@11.9.0': - resolution: {integrity: sha512-fl7W5imnSmmgXIqMQFZ/rPaVvk9OkKF8/anqHZE3XEDfWcn3BlWGndyOEas/JN7u2BXWYjs63DJZ3rnG6WOhLA==} + '@pnpm/linux-x64@11.11.0': + resolution: {integrity: sha512-rwMbNJR+PstRu+ymWoApei1CWrAnsnW3tm+3H8qOxbp8duiaj6u7DxlMzhKbVpFwylxcJdeGwZ5tReBFOVpsdw==} cpu: [x64] os: [linux] - '@pnpm/linuxstatic-arm64@11.9.0': - resolution: {integrity: sha512-fif8xbnzVEAIlvaU4yIgWKXeXYb4Kj6WMEl/KvM2x1Rp3AKAjBW/53SGzxO4cZP9doAqUzOIpMRGFVbHGeMDRw==} + '@pnpm/linuxstatic-arm64@11.11.0': + resolution: {integrity: sha512-OcmrMw1hxNee7KTBpE0yToTZziH/SCmJWwjB7YN2NmsxZVPKM2vtOWFdZufW0YN5JffKMaHbPZ2ulgYBM5qAQw==} cpu: [arm64] os: [linux] libc: [musl] - '@pnpm/linuxstatic-x64@11.9.0': - resolution: {integrity: sha512-9dKu3QdShqOpnWrjW9owARpIJeP0ul8UgIIbBUv8VDGEYibt+g49zQsNaD7SdMD3WeRSjExPQ1zIhplzr5cwvQ==} + '@pnpm/linuxstatic-x64@11.11.0': + resolution: {integrity: sha512-pWeAYeS+PPah6mXcJbOr+nPwEpyQau4iPcsqNUhTK9G5/qpG5dU1m8oHeIH83PuUUvvdJddLE1PXUkGM+QCI6g==} cpu: [x64] os: [linux] libc: [musl] - '@pnpm/macos-arm64@11.9.0': - resolution: {integrity: sha512-MWzBTgeI5p3odjdVltYvFXaSWAjF2Xk5YaxiP/u2RmW8N6PHsLyIyj37Ds992CrXgFM8fO1RpvsEirozhsD6KA==} + '@pnpm/macos-arm64@11.11.0': + resolution: {integrity: sha512-62P8Pe4yvMkdl2nxswCjM5X835i9Judk68CZvv8OnWsm7CVCEZ61SEBnNTpjJ4kXzdayFjRpPwE/jfJSvPaWww==} cpu: [arm64] os: [darwin] - '@pnpm/win-arm64@11.9.0': - resolution: {integrity: sha512-u/QxEcbKJZxC1t3zUYCZiHzu7TaZ/iXc6EGZoQjkeVT0LXmEVR6ypcK3ByjhIqbxQ3HGX75gnpIpR+VjRjubfw==} + '@pnpm/win-arm64@11.11.0': + resolution: {integrity: sha512-yYAx+A1oT+mSgrIzysuAvdYYMYEfAa+z3/UCGY3L5VC9oabs9qi71LbTan1cDui/AadrIvD177k9mV4V38KAJg==} cpu: [arm64] os: [win32] - '@pnpm/win-x64@11.9.0': - resolution: {integrity: sha512-HqJVHmZG5UKfLi38AjMl2azVmq87TlWRhwSW+f4q9LLaZeAkFyIZ7LY/pN6mh4VlH9yWj90C+tsb1yKiy7OKnw==} + '@pnpm/win-x64@11.11.0': + resolution: {integrity: sha512-ehTuyM5Rrp4ye0SdtAJkUS+9ykfwdDY9SPqiK3+bxXPx3LD5aeDw2EBksTh/xRTgqFdYmoFR86cABFt9yBr0GA==} cpu: [x64] os: [win32] @@ -116,45 +116,45 @@ packages: resolution: {integrity: sha512-Btj2BOOO83o3WyH59e8MgXsxEQVcarkUOpEYrubB0urwnN10yQ364rsiByU11nZlqWYZm05i/of7io4mzihBtQ==} engines: {node: '>=8'} - pnpm@11.9.0: - resolution: {integrity: sha512-vWgtXQP+Ul73yf1ngMaITR51asTJyf4AxTh4KCQxDc+Q493E9Tg18G3669UIXkGFXgvLs7YN4qxburieUDbwOw==} + pnpm@11.11.0: + resolution: {integrity: sha512-RGP2X9gO2A1pvB1L8WPulPYFxzgPwxi7Wy6+FfjNEtScUaTVnpUbQB52TTtsp1HL9RvFDtcAGmvLSTXmhMNIgg==} engines: {node: '>=22.13'} hasBin: true snapshots: - '@pnpm/exe@11.9.0': + '@pnpm/exe@11.11.0': dependencies: '@reflink/reflink': 0.1.19 detect-libc: 2.1.2 optionalDependencies: - '@pnpm/linux-arm64': 11.9.0 - '@pnpm/linux-x64': 11.9.0 - '@pnpm/linuxstatic-arm64': 11.9.0 - '@pnpm/linuxstatic-x64': 11.9.0 - '@pnpm/macos-arm64': 11.9.0 - '@pnpm/win-arm64': 11.9.0 - '@pnpm/win-x64': 11.9.0 + '@pnpm/linux-arm64': 11.11.0 + '@pnpm/linux-x64': 11.11.0 + '@pnpm/linuxstatic-arm64': 11.11.0 + '@pnpm/linuxstatic-x64': 11.11.0 + '@pnpm/macos-arm64': 11.11.0 + '@pnpm/win-arm64': 11.11.0 + '@pnpm/win-x64': 11.11.0 - '@pnpm/linux-arm64@11.9.0': + '@pnpm/linux-arm64@11.11.0': optional: true - '@pnpm/linux-x64@11.9.0': + '@pnpm/linux-x64@11.11.0': optional: true - '@pnpm/linuxstatic-arm64@11.9.0': + '@pnpm/linuxstatic-arm64@11.11.0': optional: true - '@pnpm/linuxstatic-x64@11.9.0': + '@pnpm/linuxstatic-x64@11.11.0': optional: true - '@pnpm/macos-arm64@11.9.0': + '@pnpm/macos-arm64@11.11.0': optional: true - '@pnpm/win-arm64@11.9.0': + '@pnpm/win-arm64@11.11.0': optional: true - '@pnpm/win-x64@11.9.0': + '@pnpm/win-x64@11.11.0': optional: true '@reflink/reflink-darwin-arm64@0.1.19': @@ -194,7 +194,7 @@ snapshots: detect-libc@2.1.2: {} - pnpm@11.9.0: {} + pnpm@11.11.0: {} --- lockfileVersion: '9.0' @@ -289,9 +289,15 @@ 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) + debug: + specifier: 'catalog:' + version: 4.4.3(supports-color@8.1.1) immutable: specifier: 'catalog:' version: 5.1.6 diff --git a/server/src/controllers/federated_controller.ts b/server/src/controllers/federated_controller.ts index d4d0fdfd2..db5d9170b 100644 --- a/server/src/controllers/federated_controller.ts +++ b/server/src/controllers/federated_controller.ts @@ -3,7 +3,7 @@ import type WebSocket from "ws"; import { v4 as randomUUID } from "uuid"; import * as msgpack from "@msgpack/msgpack"; -import type { DataType, Task, Encoded } from "@epfml/discojs"; +import type { DataType, Task, Encoded, NodeID } from "@epfml/discojs"; import { mtype, federatedMessages, @@ -22,11 +22,16 @@ export class FederatedController extends TrainingController< D, "federated" > { + /** + * WebSockets of clients whose update was accepted for the current round. + * They receive the resulting global weights when aggregation completes. + */ + #pendingUpdateRecipients = new Map(); /** * Aggregators for each hosted task. By default the server waits for 100% of the nodes to send their contributions before aggregating the updates */ - #aggregator = new 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 @@ -40,11 +45,70 @@ export class FederatedController extends TrainingController< ) { super(task); this.#latestGlobalWeights = this.initialWeights; + } + + /** + * Creates an aggregator and registers the handler that caches and broadcasts + * the global weights produced at the end of each aggregation round. + */ + #makeAggregator(): MeanAggregator { + const aggregator = new MeanAggregator(undefined, 1, "relative"); + + aggregator.on("aggregation", async (weightUpdate) => { + try { + const payload = await weightsEncode(weightUpdate); + const recipients = this.#pendingUpdateRecipients; + this.#pendingUpdateRecipients = new Map(); + + debug( + "round %o aggregate payload byteLength=%d", + aggregator.round, + payload.byteLength, + ); + this.#latestGlobalWeights = payload; - // 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 weightsEncode(weightUpdate); + const msg: federatedMessages.ReceiveServerPayload = { + type: MessageTypes.ReceiveServerPayload, + round: aggregator.round, + payload, + nbOfParticipants: this.connections.size, + }; + const encodedMsg = msgpack.encode(msg); + + recipients.forEach((recipientWs, recipientId) => { + try { + debug( + "Sending global weights for round %o to client [%s]", + aggregator.round, + recipientId.slice(0, 4), + ); + recipientWs.send(encodedMsg); + debug( + "Aggregated payload sent to client [%s] for round %o", + recipientId.slice(0, 4), + aggregator.round, + ); + } catch (err) { + debug( + "Failed to send global weights for round %o to client [%s]: %o", + aggregator.round, + recipientId.slice(0, 4), + err, + ); + } + }); + } catch (err) { + debug( + "Failed to serialize or encode weights for round %o: %o", + aggregator.round, + err, + ); + } finally { + weightUpdate.dispose(); + } }); + + return aggregator; } /** @@ -68,6 +132,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); @@ -92,7 +160,10 @@ export class FederatedController extends TrainingController< 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, }; @@ -107,30 +178,32 @@ 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 = weightsDecode(payload); - - // Create a callback to send the aggregated weight to the client - // when enough contributions are received - this.#aggregator.once("aggregation", async (weightUpdate) => { + let added = false; + try { + // Add the contribution debug( - "Sending global weights for round %o to client [%s]", - this.#aggregator.round, + "Adding contribution from client [%s] to aggregator for round %d", shortId, + round, ); - const msg: federatedMessages.ReceiveServerPayload = { - type: MessageTypes.ReceiveServerPayload, - round: this.#aggregator.round, // send the current round number after aggregation - payload: await weightsEncode(weightUpdate), - nbOfParticipants: this.connections.size, - }; - ws.send(msgpack.encode(msg)); - }); - // Add the contribution - this.#aggregator.add(clientId, weights, round); - debug( - `Successfully added contribution from client [%s] for round ${round}`, - shortId, - ); + 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 // the server answers with the current round and last global model update @@ -159,13 +232,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 MeanAggregator(undefined, 1, "relative"); + this.#pendingUpdateRecipients.clear(); + this.#aggregator = this.#makeAggregator(); this.#latestGlobalWeights = this.initialWeights; } diff --git a/server/src/routes/training_router.ts b/server/src/routes/training_router.ts index 1c75387ba..83984b02c 100644 --- a/server/src/routes/training_router.ts +++ b/server/src/routes/training_router.ts @@ -2,6 +2,7 @@ import express from "express"; import type expressWS from "express-ws"; import type { Task, DataType, Network, Encoded } from "@epfml/discojs"; import { modelDecode, weightsEncode } from "@epfml/discojs"; +import createDebug from "debug"; import type { TaskSet } from "../task_set.js"; import type { TrainingController } from "../controllers/index.js"; @@ -10,6 +11,8 @@ import { DecentralizedController, } from "../controllers/index.js"; +const debug = createDebug("server:routes:training"); + /** * The TrainingRouter handles client requests related the federated * and decentralized training. @@ -55,8 +58,17 @@ export class TrainingRouter> { // The federated controller takes the initial model weights at initialization // so that it can send it to new clients - const model = modelDecode(encodedModel); - const encodedWeights = await weightsEncode((await model).weights); + const model = await modelDecode(encodedModel); + const weights = model.weights; + let encodedWeights: Encoded; + try { + encodedWeights = await weightsEncode(weights); + } catch (err) { + debug("Failed to encode initial weights for task %s: %o", task.id, err); + throw err; + } finally { + model[Symbol.dispose](); + } taskController = new FederatedController(t, encodedWeights); } else { const t = task as Task; diff --git a/server/src/server.ts b/server/src/server.ts index 11fa58ba9..2d1b3303b 100644 --- a/server/src/server.ts +++ b/server/src/server.ts @@ -59,6 +59,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; diff --git a/server/tests/e2e/decentralized.spec.ts b/server/tests/e2e/decentralized.spec.ts index 584f1f20d..15e4051dd 100644 --- a/server/tests/e2e/decentralized.spec.ts +++ b/server/tests/e2e/decentralized.spec.ts @@ -205,10 +205,15 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { * 3. set status to "connecting to peers" and establish the connections * 4. set status to "updating model" and exchange weight updates * - * Given this, it is important to note that calling disco.trainByRound().next() - * for the first time will perform a) and then b) where it stops and yields the round logs. - * Thus, c) isn't called and the weight sharing is not performed during this call to next(). - * Calling next() again will then run c), as well as a) and b) again. + * Given this, it is important to note that a single call to + * disco.trainByRound().next() performs a full round: a), b) and c). + * It only resolves once the peers exchanged their weight updates, so when + * a peer is alone (minNbOfParticipants isn't met) the call stays pending + * until another participant joins and the round completes. The test + * therefore holds the pending next() promises and choreographs through + * the status and participants events instead of awaiting next() right + * away. Note that RoundLogs.participants is the count seen at the end of + * local training, before the weight exchange. * * In this test the timeline is: * - User 1 joins the task by themselves @@ -231,25 +236,19 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }); const generatorUser1 = discoUser1.trainByRound(dataset); - // Have User 1 join the task and train locally for one round - const logUser1Round1 = await generatorUser1.next(); - expect(logUser1Round1.done).to.be.false; + // Have User 1 join the task and train locally. The first next() call + // runs a), b) and c): the round can't complete while User 1 is alone so + // the promise stays pending in c) + const logUser1Round1Promise = generatorUser1.next(); // User 1 did a) and b) so their status should be Training expect(await statusUser1.next()).equal("local training"); expect(await nbParticipantsUser1.next()).equal(1); - - if (logUser1Round1.done) - throw new Error("User 1 finished training at the 1st round"); - // participant list not updated yet (updated at step c)) - expect(logUser1Round1.value.participants).equal(1); - - // Calling next() a 2nd time makes User 1 go to c) where the peer should - // stay stuck awaiting until another participant joins - const logUser1Round2Promise = generatorUser1.next(); + // User 1 then reaches c) where it is ready to share its weights + // but has to wait for more participants expect(await statusUser1.next()).equal( "waiting for peers to share weights", - ); // ready to share - expect(await statusUser1.next()).equal("not enough participants"); // but has to wait for more participants + ); + expect(await statusUser1.next()).equal("not enough participants"); /* USER 2 JOINS */ @@ -265,48 +264,72 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { const generatorUser2 = discoUser2.trainByRound(dataset); // Have User 2 join the task and train for one round - const logUser2Round1 = await generatorUser2.next(); - expect(logUser2Round1.done).to.be.false; - if (logUser2Round1.done) - throw new Error("User 2 finished training at the 1st round"); - // round payload should contain the number of participants - expect(logUser2Round1.value.participants).equal(2); + const logUser2Round1Promise = generatorUser2.next(); + // User 2 connects to the server which triggers the participant event expect(await nbParticipantsUser2.next()).equal(2); - // Receive the EnoughParticipants message with the participants - expect(await nbParticipantsUser1.next()).equal(2); // User 2 did a) and b) expect(await statusUser2.next()).equal("local training"); - // User 1 is still in c) now waiting for user 2 to be ready to exchange weight updates + // User 1 receives the EnoughParticipants message with the participants + // and its previous status is restored + expect(await nbParticipantsUser1.next()).equal(2); expect(await statusUser1.next()).equal( "waiting for peers to share weights", ); + // User 2 finishes training and is ready to share its weights too + expect(await statusUser2.next()).equal( + "waiting for peers to share weights", + ); + + /* ROUND 1 COMPLETES */ + + // The server answers with the round's peers list, peers exchange their + // updates and both pending next() calls resolve + const logUser1Round1 = await logUser1Round1Promise; + const logUser2Round1 = await logUser2Round1Promise; + expect(logUser1Round1.done).to.be.false; + expect(logUser2Round1.done).to.be.false; + if (logUser1Round1.done || logUser2Round1.done) + throw new Error("User 1 or 2 finished training at the 1st round"); + // User 1 finished training alone, User 2 with both present + expect(logUser1Round1.value.participants).equal(1); + expect(logUser2Round1.value.participants).equal(2); + // Receiving the peers list updates the participants + expect(await nbParticipantsUser1.next()).equal(2); + expect(await nbParticipantsUser2.next()).equal(2); + expect(await statusUser1.next()).equal("connecting to peers"); + expect(await statusUser1.next()).equal("updating model"); + expect(await statusUser2.next()).equal("connecting to peers"); + expect(await statusUser2.next()).equal("updating model"); /* ROUND 2 */ - // The server should answer with the round's peers list. - // Peers then exchange updates and then start training locally with the new weights - const logUser2Round2 = await generatorUser2.next(); - const logUser1Round2 = await logUser1Round2Promise; // the promise can resolve now + // Both users are present so the round runs a), b) and c) to completion + const [logUser1Round2, logUser2Round2] = await Promise.all([ + generatorUser1.next(), + generatorUser2.next(), + ]); expect(logUser1Round2.done).to.be.false; expect(logUser2Round2.done).to.be.false; if (logUser1Round2.done || logUser2Round2.done) throw new Error("User 1 or 2 finished training at the 2nd round"); - // nb of participants should now be updated expect(logUser1Round2.value.participants).equal(2); expect(logUser2Round2.value.participants).equal(2); - expect(await nbParticipantsUser2.next()).equal(2); expect(await nbParticipantsUser1.next()).equal(2); - // User 1 and 2 did c), a) and b) - expect(await statusUser1.next()).equal("connecting to peers"); - expect(await statusUser1.next()).equal("updating model"); // second to last + expect(await nbParticipantsUser2.next()).equal(2); + // User 1 and 2 did a), b) and c) expect(await statusUser1.next()).equal("local training"); + expect(await statusUser1.next()).equal( + "waiting for peers to share weights", + ); + expect(await statusUser1.next()).equal("connecting to peers"); + expect(await statusUser1.next()).equal("updating model"); + expect(await statusUser2.next()).equal("local training"); expect(await statusUser2.next()).equal( "waiting for peers to share weights", ); expect(await statusUser2.next()).equal("connecting to peers"); expect(await statusUser2.next()).equal("updating model"); - expect(await statusUser2.next()).equal("local training"); /* USER 1 LEAVES */ @@ -317,10 +340,10 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { expect(await nbParticipantsUser2.next()).equal(1); // server notifies user 2 to wait expect(await statusUser2.next()).equal("not enough participants"); - // Make user 2 go to c) + // Make User 2 start round 3: it trains, then waits in c) for another + // participant, so the promise stays pending const logUser2Round3Promise = generatorUser2.next(); - // await new Promise((res, _) => setTimeout(res, statusUpdateTime)) // Wait some time for the status to update - // starts c) and waits for user 3 to join + expect(await statusUser2.next()).equal("local training"); expect(await statusUser2.next()).equal( "waiting for peers to share weights", ); @@ -340,49 +363,45 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }); const generatorUser3 = discoUser3.trainByRound(dataset); - // User 3 joins mid-training and trains one local round - const logUser3Round1 = await generatorUser3.next(); - expect(logUser3Round1.done).to.be.false; - if (logUser3Round1.done) - throw new Error("User 3 finished training at the 1st round"); - expect(logUser3Round1.value.participants).equal(2); + // User 3 joins mid-training and trains one local round; the round can + // only complete once User 3 shares its update so hold the promise + const logUser3Round1Promise = generatorUser3.next(); expect(await nbParticipantsUser3.next()).equal(2); - // User 2 receives the EnoughParticipants message - // User 2 is still in c) waiting for user 3 to share their local update - expect(await nbParticipantsUser2.next()).equal(2); - // User 3 did a) and b) expect(await statusUser3.next()).equal("local training"); - // User 2 is still in c) waiting for user 3 to be ready to exchange waits + // User 2 receives the EnoughParticipants message and its previous + // status is restored, waiting for user 3 to be ready to exchange weights + expect(await nbParticipantsUser2.next()).equal(2); expect(await statusUser2.next()).equal( "waiting for peers to share weights", ); + // User 3 finishes training and notifies the server that they are ready + expect(await statusUser3.next()).equal( + "waiting for peers to share weights", + ); - /* ROUND 3 */ + /* ROUND 3 COMPLETES */ - // User 3 notifies the server that they are ready to exchange waits - // then user 2 and 3 exchange weight updates - const logUser3Round3 = await generatorUser3.next(); - const logUser2Round3 = await logUser2Round3Promise; // the promise can resolve now - if (logUser3Round3.done || logUser2Round3.done) - throw new Error("User 1 or 2 finished training at the 3rd round"); + // user 2 and 3 exchange weight updates and both rounds resolve + const logUser2Round3 = await logUser2Round3Promise; + const logUser3Round1 = await logUser3Round1Promise; + expect(logUser2Round3.done).to.be.false; + expect(logUser3Round1.done).to.be.false; + if (logUser2Round3.done || logUser3Round1.done) + throw new Error("User 2 or 3 finished training at the 3rd round"); - expect(logUser2Round3.value.participants).equal(2); - expect(logUser3Round3.value.participants).equal(2); - expect(await nbParticipantsUser3.next()).equal(2); + // User 2 finished training alone, User 3 with both present + expect(logUser2Round3.value.participants).equal(1); + expect(logUser3Round1.value.participants).equal(2); expect(await nbParticipantsUser2.next()).equal(2); + expect(await nbParticipantsUser3.next()).equal(2); - // both user 2 and 3 did c), a) and are now in b) + // both user 2 and 3 finish c) expect(await statusUser2.next()).equal("connecting to peers"); expect(await statusUser2.next()).equal("updating model"); - expect(await statusUser2.next()).equal("local training"); - expect(await statusUser3.next()).equal( - "waiting for peers to share weights", - ); expect(await statusUser3.next()).equal("connecting to peers"); expect(await statusUser3.next()).equal("updating model"); - expect(await statusUser3.next()).equal("local training"); /* USER 2 AND 3 LEAVE */ @@ -404,13 +423,17 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { const dataset = await datasets.loadLusCOVID(); /** - * The timeline is: - * - User 1 joins the task by themselves and trains locally - * - User 2 joins while User 1 is still training - * - User 1 is done training and waits for User 2 to share its weights + * A call to trainByRound().next() runs a full round: a), b) and c), + * so User 1's first round only resolves once User 2 joined and the + * round was aggregated. The timeline is: + * - User 1 joins the task by themselves, trains locally and waits in c) + * to share its weights ("not enough participants" is expected there, + * User 1 really is alone) + * - User 2 joins and trains while User 1 is ready to share * - * User 1 has to wait for User 2 to be ready but shouldn't be told that - * participants are missing: User 2 is here, only still training. + * User 1 has to wait for User 2 to be ready but, once User 2 joined, + * shouldn't be told that participants are missing: User 2 is here, only + * still training. */ /* USER 1 JOINS */ @@ -422,31 +445,45 @@ describe("end-to-end decentralized", { timeout: 50_000 }, () => { }); const generatorUser1 = discoUser1.trainByRound(dataset); - await generatorUser1.next(); // a) and b) + // a), b) and c): User 1 trains then waits in c) until User 2 joins + // and the round completes, so the promise stays pending + const logUser1Round1Promise = generatorUser1.next(); expect(await statusUser1.next()).equal("local training"); + expect(await statusUser1.next()).equal( + "waiting for peers to share weights", + ); + // User 1 is genuinely alone at this point + expect(await statusUser1.next()).equal("not enough participants"); - /* USER 2 JOINS, WHILE USER 1 IS STILL TRAINING */ + /* USER 2 JOINS, WHILE USER 1 IS READY TO SHARE ITS WEIGHTS */ const discoUser2 = new Disco(task, url, { preprocessOnce: true }); + const statusUser2 = new Queue(); + discoUser2.on("status", (status) => { + statusUser2.put(status); + }); const generatorUser2 = discoUser2.trainByRound(dataset); - await generatorUser2.next(); // a) and b) - - // there are enough participants now, User 1 keeps on training - expect(await statusUser1.next()).equal("local training"); + const logUser2Round1Promise = generatorUser2.next(); // a), b) and c) - /* USER 1 IS DONE TRAINING */ - - const logUser1Round2 = generatorUser1.next(); // c) + // There are enough participants now: User 1's status is restored while + // User 2 trains. User 1 waits for User 2 to be ready but should NOT + // report missing participants: the next status updates must be + // "connecting to peers", without any "not enough participants" expect(await statusUser1.next()).equal( "waiting for peers to share weights", ); + expect(await statusUser2.next()).equal("local training"); + expect(await statusUser2.next()).equal( + "waiting for peers to share weights", + ); - /* USER 2 IS DONE TRAINING TOO */ + /* USER 2 IS DONE TRAINING, PEERS EXCHANGE THEIR UPDATES */ - await generatorUser2.next(); - await logUser1Round2; + await Promise.all([logUser1Round1Promise, logUser2Round1Promise]); expect(await statusUser1.next()).equal("connecting to peers"); expect(await statusUser1.next()).equal("updating model"); + expect(await statusUser2.next()).equal("connecting to peers"); + expect(await statusUser2.next()).equal("updating model"); await discoUser1.close(); await discoUser2.close(); diff --git a/server/tests/e2e/federated.spec.ts b/server/tests/e2e/federated.spec.ts index 072aa2372..fe55f8933 100644 --- a/server/tests/e2e/federated.spec.ts +++ b/server/tests/e2e/federated.spec.ts @@ -10,7 +10,7 @@ import type { WeightsContainer, ModelCard, } from "@epfml/discojs"; -import { Disco, defaultTasks, defaultModels } from "@epfml/discojs"; +import { Disco, defaultTasks, defaultModels, GPT } from "@epfml/discojs"; import { List } from "immutable"; import { assert, afterEach, describe, expect, it } from "vitest"; import { Server } from "../../src/index.js"; @@ -176,7 +176,17 @@ describe("end-to-end federated", () => { ...defaultTasks.wikitext, getTask: () => Promise.resolve(task), }; - const url = await startServer(defaultModels.Wikitext, taskProvider); + const wikitextModelCard = { + ...defaultModels.Wikitext, + getModel: () => + Promise.resolve( + new GPT({ + contextLength: task.trainingInformation.contextLength, + maxIter: 10, + }), + ), + }; + const url = await startServer(wikitextModelCard, taskProvider); const dataset = datasets.loadWikitext(); const [r1, r2] = await Promise.all([ @@ -209,11 +219,13 @@ describe("end-to-end federated", () => { * c) onRoundEndCommunication which sends the local update and * receives the global weights while emitting the status UPDATE * - * Given this, it is important to note that calling disco.trainByRound().next() - * for the first time will perform a) and then b) where it stops and yields the round logs. - * Thus, c) isn't done and the model aggregation by the server is not performed during this first call to next(). - * - * Calling next() again will then do c), and back to a) and b). + * Given this, it is important to note that a single call to + * disco.trainByRound().next() performs a full round: a), b) and c). + * It only resolves once the server aggregated the round, so when a client + * is alone (minNbOfParticipants isn't met) the call stays pending until + * another participant joins and the round completes. Tests therefore hold + * the pending next() promise and choreograph through the status and + * participants events instead of awaiting next() right away. * * In this test the timeline is: * - User 1 joins the task by themselves @@ -233,14 +245,11 @@ describe("end-to-end federated", () => { ); const generatorUser1 = discoUser1.trainByRound(dataset); - // Have User 1 join the task and train locally for one round - await generatorUser1.next(); + // Have User 1 join the task and train locally. The round can't complete + // while User 1 is alone so the promise stays pending in c) + const logUser1Round1Promise = generatorUser1.next(); expect(await statusUser1.next()).equal("local training"); expect(await nbParticipantsUser1.next()).equal(1); - - // Calling next() a 2nd time makes User 1 go to c) where the client should - // stay stuck awaiting until another participant joins - const logUser1Round2Promise = generatorUser1.next(); expect(await statusUser1.next()).equal("not enough participants"); // Create User 2 @@ -254,34 +263,34 @@ describe("end-to-end federated", () => { const generatorUser2 = discoUser2.trainByRound(dataset); // Have User 2 join the task and train for one round - await generatorUser2.next(); - // User 2 did a) and b) - expect(await statusUser1.next()).equal("local training"); - expect(await statusUser2.next()).equal("local training"); - // User 1 is still in c) now waiting for user 2 to share their local update - // and for the server to aggregate the local updates - expect(await statusUser1.next()).equal("updating model"); + const logUser2Round1Promise = generatorUser2.next(); // User 2 connects to the server which triggers the participant event expect(await nbParticipantsUser2.next()).equal(2); - // Receive the EnoughParticipants message with the participants + expect(await statusUser2.next()).equal("local training"); + // User 1 receives the EnoughParticipants message with the participants, + // its previous status is restored and it proceeds to share its update expect(await nbParticipantsUser1.next()).equal(2); + expect(await statusUser1.next()).equal("local training"); + expect(await statusUser1.next()).equal("updating model"); + // User 2 finishes training and shares its update too + expect(await statusUser2.next()).equal("updating model"); - // Proceed with round 2 + // The server aggregates the round and answers with the new global weights + // along with the participants, resolving both pending next() calls + await Promise.all([logUser1Round1Promise, logUser2Round1Promise]); + expect(await nbParticipantsUser1.next()).equal(2); + expect(await nbParticipantsUser2.next()).equal(2); - // the server should answer with the new global weights - // and users should train locally on the new weights - await Promise.all([logUser1Round2Promise, generatorUser2.next()]); - // User 1 and 2 did c), a) and b) - expect(await statusUser2.next()).equal("updating model"); + // Proceed with round 2, both users are present so the round completes + await Promise.all([generatorUser1.next(), generatorUser2.next()]); + // User 1 and 2 did a), b) and c) expect(await statusUser1.next()).equal("local training"); + expect(await statusUser1.next()).equal("updating model"); expect(await statusUser2.next()).equal("local training"); + expect(await statusUser2.next()).equal("updating model"); // Receive the server payload during c) along with the participants - expect(await nbParticipantsUser2.next()).equal(2); expect(await nbParticipantsUser1.next()).equal(2); - - // Make user 2 go to c) - const logUser2Round3Promise = generatorUser2.next(); - expect(await statusUser2.next()).equal("updating model"); + expect(await nbParticipantsUser2.next()).equal(2); // Have user 1 quit the session await discoUser1.close(); @@ -289,6 +298,12 @@ describe("end-to-end federated", () => { expect(await statusUser2.next()).equal("not enough participants"); expect(await nbParticipantsUser2.next()).equal(1); + // Make User 2 start round 3, it trains and then waits in c) for + // another participant + const logUser2Round3Promise = generatorUser2.next(); + expect(await statusUser2.next()).equal("local training"); + expect(await statusUser2.next()).equal("not enough participants"); + // Create User 3 const discoUser3 = new Disco(task, url, { preprocessOnce: true }); const statusUser3 = new Queue(); @@ -300,27 +315,24 @@ describe("end-to-end federated", () => { const generatorUser3 = discoUser3.trainByRound(dataset); // User 3 joins mid-training and trains one local round - await generatorUser3.next(); - expect(await statusUser3.next()).equal("local training"); + const logUser3Round1Promise = generatorUser3.next(); expect(await nbParticipantsUser3.next()).equal(2); + expect(await statusUser3.next()).equal("local training"); - // User 2 is still in c) waiting for user 3 to share their local update - // and for the server to aggregate the local updates - expect(await statusUser2.next()).equal("updating model"); - // User 2 receives the EnoughParticipants message + // User 2 receives the EnoughParticipants message, its previous status + // is restored and it proceeds to share its update expect(await nbParticipantsUser2.next()).equal(2); - - // User 3 sends their weights to the server - await Promise.all([logUser2Round3Promise, generatorUser3.next()]); + expect(await statusUser2.next()).equal("local training"); + expect(await statusUser2.next()).equal("updating model"); + // User 3 finishes training and sends their weights to the server expect(await statusUser3.next()).equal("updating model"); - // the server should accept user 3's weights (should not be outdated) and aggregate the global weights - // both user 2 and 3 did c), a) and are now in b) - expect(await statusUser2.next()).equal("local training"); - expect(await statusUser3.next()).equal("local training"); + // the server should accept user 3's weights (should not be outdated) + // and aggregate the global weights, resolving both rounds + await Promise.all([logUser2Round3Promise, logUser3Round1Promise]); // User 2 and 3 finish c) - expect(await nbParticipantsUser3.next()).equal(2); expect(await nbParticipantsUser2.next()).equal(2); + expect(await nbParticipantsUser3.next()).equal(2); await discoUser2.close(); expect(await statusUser3.next()).equal("not enough participants");