@epfml/discojs 2.1.2-p20240627125649.0 → 2.2.0

package/dist/client/decentralized/peer.d.ts

@@ -1,4 +1,3 @@
-/// <reference types="node" resolution-mode="require"/>
 import type { NodeID } from '../types.js';
 export type SignalData = {
     type: 'answer' | 'offer' | 'pranswer' | 'rollback';

package/dist/default_tasks/wikitext.js

@@ -20,7 +20,7 @@ export const wikitext = {
                 modelID: 'llm-raw-model',
                 preprocessingFunctions: [data.TextPreprocessing.Tokenize, data.TextPreprocessing.LeftPadding],
                 scheme: 'federated',
-                epochs: 5,
+                epochs: 6,
                 // Unused by wikitext because data already comes split
                 // But if set to 0 then the webapp doesn't display the validation metrics
                 validationSplit: 0.1,

package/dist/index.d.ts

@@ -10,8 +10,9 @@ export { Logger, ConsoleLogger } from './logging/index.js';
 export { Memory, type ModelInfo, type Path, type ModelSource, Empty as EmptyMemory } from './memory/index.js';
 export { Disco, RoundLogs } from './training/index.js';
 export { Validator } from './validation/index.js';
-export { Model, EpochLogs } from './models/index.js';
+export { Model, BatchLogs, EpochLogs, ValidationMetrics } from './models/index.js';
 export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
 export * from './types.js';
+export * as async_iterator from "./utils/async_iterator.js";

package/dist/index.js

@@ -10,8 +10,9 @@ export { ConsoleLogger } from './logging/index.js';
 export { Memory, Empty as EmptyMemory } from './memory/index.js';
 export { Disco } from './training/index.js';
 export { Validator } from './validation/index.js';
-export { Model } from './models/index.js';
+export { Model, EpochLogs } from './models/index.js';
 export * as models from './models/index.js';
 export * from './task/index.js';
 export * as defaultTasks from './default_tasks/index.js';
 export * from './types.js';
+export * as async_iterator from "./utils/async_iterator.js";

package/dist/models/gpt/evaluate.d.ts

@@ -3,5 +3,5 @@ interface DataPoint extends tf.TensorContainerObject {
     xs: tf.Tensor2D;
     ys: tf.Tensor3D;
 }
-export default function evaluate(model: tf.LayersModel, dataset: tf.data.Dataset<DataPoint>, maxEvalBatches: number): Promise<Record<'acc' | 'val_acc' | 'val_loss' | 'val_perplexity', number>>;
+export default function evaluate(model: tf.LayersModel, dataset: tf.data.Dataset<DataPoint>, maxEvalBatches: number): Promise<Record<'val_acc' | 'val_loss' | 'val_perplexity', number>>;
 export {};

package/dist/models/gpt/evaluate.js

@@ -38,7 +38,6 @@ export default async function evaluate(model, dataset, maxEvalBatches) {
     return {
         val_loss: loss,
         val_perplexity: Math.exp(loss),
-        acc: acc[0] / acc[1],
         val_acc: acc[0] / acc[1]
     };
 }

package/dist/models/gpt/index.d.ts

@@ -5,14 +5,15 @@ import * as tf from '@tensorflow/tfjs';
 import { PreTrainedTokenizer } from '@xenova/transformers';
 import { WeightsContainer } from '../../index.js';
 import type { Dataset } from '../../dataset/index.js';
-import { Model } from '../model.js';
-import type { EpochLogs, Prediction, Sample } from '../model.js';
-import type { GPTConfig } from './config.js';
+import { BatchLogs, Model, EpochLogs } from "../index.js";
+import type { Prediction, Sample } from '../model.js';
+import { type GPTConfig } from './config.js';
 export type GPTSerialization = {
     weights: WeightsContainer;
     config?: GPTConfig;
 };
 export declare class GPT extends Model {
+    #private;
     private readonly model;
     constructor(partialConfig?: GPTConfig, layersModel?: tf.LayersModel);
     /**
@@ -24,7 +25,7 @@ export declare class GPT extends Model {
      * @param epochs the number of passes of the training dataset
      * @param tracker
      */
-    train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs, void>;
+    train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
     predict(input: Sample): Promise<Prediction>;
     generate(input: string, tokenizer: PreTrainedTokenizer, newTokens?: number): Promise<string>;
     get config(): Required<GPTConfig>;

package/dist/models/gpt/index.js

@@ -1,14 +1,20 @@
 /**
  * this code is taken from gpt-tfjs with modifications from @peacefulotter and @lukemovement
  **/
+import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../../index.js';
-import { Model } from '../model.js';
+import { Model, EpochLogs } from "../index.js";
 import { GPTForCausalLM } from './model.js';
+import { DEFAULT_CONFIG } from './config.js';
+import evaluate from './evaluate.js';
+import { List } from 'immutable';
 export class GPT extends Model {
     model;
+    #maxBatchCount;
     constructor(partialConfig, layersModel) {
         super();
         this.model = new GPTForCausalLM(partialConfig, layersModel);
+        this.#maxBatchCount = partialConfig?.maxIter ?? DEFAULT_CONFIG.maxIter;
     }
     /**
      * The GPT train methods wraps the model.fitDataset call in a for loop to act as a generator (of logs)
@@ -19,45 +25,65 @@ export class GPT extends Model {
      * @param epochs the number of passes of the training dataset
      * @param tracker
      */
-    async *train(trainingData, validationData, epochs = 1) {
+    async *train(trainingData, validationData) {
         this.model.compile();
+        const batches = await trainingData.iterator(); // tf.LazyIterator isn't an AsyncGenerator
+        let batchesLogs = List();
+        for (let batchNumber = 0; batchNumber < this.#maxBatchCount; batchNumber++) {
+            const iteration = await batches.next();
+            if (iteration.done)
+                break;
+            const batch = iteration.value;
+            const batchLogs = await this.#runBatch(batch);
+            tf.dispose(batch);
+            yield batchLogs;
+            batchesLogs = batchesLogs.push(batchLogs);
+        }
+        const validation = validationData && (await this.#evaluate(validationData));
+        return new EpochLogs(batchesLogs, validation);
+    }
+    async #runBatch(batch) {
         let logs;
-        const trainingArgs = {
-            epochs: 1, // force fitDataset to do only one epoch because it is wrapped in a for loop
-            validationData,
-            callbacks: { onEpochEnd: (_, cur) => { logs = cur; } },
+        await this.model.fitDataset(tf.data.array([batch]), {
+            epochs: 1,
+            verbose: 0, // don't pollute
+            callbacks: {
+                onEpochEnd: (_, cur) => {
+                    logs = cur;
+                },
+            },
+        });
+        if (logs === undefined)
+            throw new Error("batch didn't gave any logs");
+        const { loss, acc: accuracy } = logs;
+        if (loss === undefined || isNaN(loss))
+            throw new Error("training loss is undefined or NaN");
+        return {
+            accuracy,
+            loss,
+            memoryUsage: tf.memory().numBytes / 1024 / 1024 / 1024,
         };
-        for (let epoch = 0; epoch < epochs; epoch++) {
-            await this.model.fitDataset(trainingData, trainingArgs);
-            if (logs === undefined) {
-                throw new Error("Epoch didn't gave any logs");
-            }
-            const { loss, val_acc, val_loss, peakMemory } = logs;
-            if (loss === undefined || isNaN(loss)) {
-                throw new Error("Training loss is undefined or nan");
-            }
-            const structuredLogs = {
-                epoch,
-                peakMemory,
-                training: {
-                    loss: logs.loss,
-                    accuracy: logs.acc
-                }
-            };
-            if (validationData !== undefined) {
-                if (val_loss === undefined || isNaN(val_loss) ||
-                    val_acc === undefined || isNaN(val_acc)) {
-                    throw new Error("Validation accuracy or loss is undefined or nan");
-                }
-                structuredLogs.validation = { accuracy: logs.val_acc, loss: logs.val_loss };
+    }
+    async #evaluate(dataset) {
+        const evaluation = await evaluate(this.model, dataset.map((t) => {
+            switch (t) {
+                case null:
+                case undefined:
+                    throw new Error("nullish value in dataset");
+                default:
+                    // TODO unsafe cast
+                    return t;
             }
-            yield structuredLogs;
-        }
+        }), this.config.maxEvalBatches);
+        return {
+            accuracy: evaluation.val_acc,
+            loss: evaluation.val_loss,
+        };
     }
     predict(input) {
         const ret = this.model.predict(input);
         if (Array.isArray(ret)) {
-            throw new Error('prediction yield many Tensors but should have only returned one');
+            throw new Error("prediction yield many Tensors but should have only returned one");
         }
         return Promise.resolve(ret);
     }

package/dist/models/gpt/model.js

@@ -37,29 +37,44 @@ class GPTModel extends tf.LayersModel {
         const evalDataset = trainingArgs.validationData;
         await callbacks.onTrainBegin?.();
         for (let epoch = 1; epoch <= trainingArgs.epochs; epoch++) {
+            let accuracyFraction = [0, 0];
             let averageLoss = 0;
-            let peakMemory = 0;
             let iteration = 1;
             const iterator = await dataset.iterator();
-            let preprocessingTime = performance.now();
             let next = await iterator.next();
-            preprocessingTime = performance.now() - preprocessingTime;
             while (next.done !== true && iteration <= this.config.maxIter) {
                 let weightUpdateTime = performance.now();
                 await callbacks.onEpochBegin?.(epoch);
                 const { xs, ys } = next.value;
-                const lossFn = () => {
+                let preprocessingTime = performance.now();
+                await Promise.all([xs.data(), ys.data()]);
+                preprocessingTime = performance.now() - preprocessingTime;
+                // TODO include as a tensor inside the model
+                const accTensor = tf.tidy(() => {
                     const logits = this.apply(xs);
-                    if (Array.isArray(logits)) {
+                    if (Array.isArray(logits))
                         throw new Error('model outputs too many tensor');
-                    }
-                    if (logits instanceof tf.SymbolicTensor) {
+                    if (logits instanceof tf.SymbolicTensor)
                         throw new Error('model outputs symbolic tensor');
-                    }
-                    return tf.losses.softmaxCrossEntropy(ys, logits);
-                };
+                    return tf.metrics.categoricalAccuracy(ys, logits);
+                });
+                const accSize = accTensor.shape.reduce((l, r) => l * r, 1);
+                const accSumTensor = accTensor.sum();
+                const accSum = await accSumTensor.array();
+                tf.dispose(accSumTensor);
+                if (typeof accSum !== 'number')
+                    throw new Error('got multiple accuracy sum');
+                accuracyFraction = [accuracyFraction[0] + accSum, accuracyFraction[1] + accSize];
+                tf.dispose([accTensor]);
                 const lossTensor = tf.tidy(() => {
-                    const { grads, value: lossTensor } = this.optimizer.computeGradients(lossFn);
+                    const { grads, value: lossTensor } = this.optimizer.computeGradients(() => {
+                        const logits = this.apply(xs);
+                        if (Array.isArray(logits))
+                            throw new Error('model outputs too many tensor');
+                        if (logits instanceof tf.SymbolicTensor)
+                            throw new Error('model outputs symbolic tensor');
+                        return tf.losses.softmaxCrossEntropy(ys, logits);
+                    });
                     const gradsClipped = clipByGlobalNormObj(grads, 1);
                     this.optimizer.applyGradients(gradsClipped);
                     return lossTensor;
@@ -75,9 +90,6 @@ class GPTModel extends tf.LayersModel {
                     console.log(iterationLogs);
                 }
                 const memory = tf.memory().numBytes / 1024 / 1024 / 1024;
-                if (memory > peakMemory) {
-                    peakMemory = memory;
-                }
                 console.log(`Epoch: ${epoch}`, `\tStep: ${iteration} / ${this.config.maxIter}`, `\tLoss: ${loss.toFixed(3)}`, `\tMemory: ${memory.toFixed(2)} GB`, `\tNumber of tensors allocated: ${tf.memory().numTensors}`, `\tPreprocessing time: ${preprocessingTime.toFixed(0)} ms`, `\tWeight update time: ${weightUpdateTime.toFixed(0)} ms`);
                 iteration++;
                 next = await iterator.next();
@@ -89,7 +101,7 @@ class GPTModel extends tf.LayersModel {
             }
             let logs = {
                 'loss': averageLoss / iteration,
-                'peakMemory': peakMemory
+                'acc': accuracyFraction[0] / accuracyFraction[1],
             };
             if (evalDataset !== undefined) {
                 logs = { ...logs, ...await evaluate(this, evalDataset, this.config.maxEvalBatches) };

package/dist/models/index.d.ts

@@ -1,4 +1,5 @@
-export { EpochLogs, Model } from './model.js';
+export { Model } from './model.js';
+export { BatchLogs, EpochLogs, ValidationMetrics } from "./logs.js";
 export { GPT } from './gpt/index.js';
 export { GPTConfig } from './gpt/config.js';
 export { TFJS } from './tfjs.js';

package/dist/models/index.js

@@ -1,4 +1,5 @@
 export { Model } from './model.js';
+export { EpochLogs } from "./logs.js";
 export { GPT } from './gpt/index.js';
 export { TFJS } from './tfjs.js';
 export { getTaskTokenizer } from './tokenizer.js';

package/dist/models/logs.d.ts

@@ -0,0 +1,17 @@
+import { List } from "immutable";
+export interface ValidationMetrics {
+    accuracy: number;
+    loss: number;
+}
+export interface BatchLogs {
+    accuracy: number;
+    loss: number;
+    memoryUsage: number;
+}
+export declare class EpochLogs {
+    readonly validation?: ValidationMetrics | undefined;
+    readonly batches: List<BatchLogs>;
+    constructor(batches: Iterable<BatchLogs>, validation?: ValidationMetrics | undefined);
+    get training(): Record<"accuracy" | "loss", number>;
+    get peakMemory(): number;
+}

package/dist/models/logs.js

@@ -0,0 +1,22 @@
+import { List } from "immutable";
+export class EpochLogs {
+    validation;
+    batches;
+    constructor(batches, validation) {
+        this.validation = validation;
+        this.batches = List(batches);
+    }
+    get training() {
+        const sum = this.batches.reduce((acc, batch) => ({
+            accuracy: acc.accuracy + batch.accuracy,
+            loss: acc.loss + batch.loss,
+        }), { loss: 0, accuracy: 0 });
+        return {
+            accuracy: sum.accuracy / this.batches.size,
+            loss: sum.loss / this.batches.size,
+        };
+    }
+    get peakMemory() {
+        return this.batches.map((batch) => batch.memoryUsage).max() ?? 0;
+    }
+}

package/dist/models/model.d.ts

@@ -1,19 +1,7 @@
-/// <reference types="node" resolution-mode="require"/>
 import type tf from "@tensorflow/tfjs";
 import type { WeightsContainer } from "../index.js";
 import type { Dataset } from "../dataset/index.js";
-export interface EpochLogs {
-    epoch: number;
-    training: {
-        loss: number;
-        accuracy?: number;
-    };
-    validation?: {
-        loss: number;
-        accuracy: number;
-    };
-    peakMemory: number;
-}
+import type { BatchLogs, EpochLogs } from "./logs.js";
 export type Prediction = tf.Tensor;
 export type Sample = tf.Tensor;
 /**
@@ -35,7 +23,7 @@ export declare abstract class Model implements Disposable {
      * @param tracker watch the various steps
      * @yields on every epoch, training can be stop by `return`ing it
      */
-    abstract train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs, void>;
+    abstract train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
     /** Predict likely values */
     abstract predict(input: Sample): Promise<Prediction>;
     /**

package/dist/models/tfjs.d.ts

@@ -1,16 +1,18 @@
 import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../index.js';
-import { Model } from './index.js';
-import type { EpochLogs, Prediction, Sample } from './model.js';
 import type { Dataset } from '../dataset/index.js';
+import { BatchLogs, EpochLogs } from './index.js';
+import { Model } from './index.js';
+import type { Prediction, Sample } from './model.js';
 /** TensorFlow JavaScript model with standard training */
 export declare class TFJS extends Model {
+    #private;
     private readonly model;
     /** Wrap the given trainable model */
     constructor(model: tf.LayersModel);
     get weights(): WeightsContainer;
     set weights(ws: WeightsContainer);
-    train(trainingData: Dataset, validationData?: Dataset, epochs?: number): AsyncGenerator<EpochLogs>;
+    train(trainingData: Dataset, validationData?: Dataset): AsyncGenerator<BatchLogs, EpochLogs>;
     predict(input: Sample): Promise<Prediction>;
     static deserialize(raw: tf.io.ModelArtifacts): Promise<Model>;
     serialize(): Promise<tf.io.ModelArtifacts>;

package/dist/models/tfjs.js

@@ -1,5 +1,7 @@
+import { List, Map } from 'immutable';
 import * as tf from '@tensorflow/tfjs';
 import { WeightsContainer } from '../index.js';
+import { EpochLogs } from './index.js';
 import { Model } from './index.js';
 /** TensorFlow JavaScript model with standard training */
 export class TFJS extends Model {
@@ -18,50 +20,71 @@ export class TFJS extends Model {
     set weights(ws) {
         this.model.setWeights(ws.weights);
     }
-    async *train(trainingData, validationData, epochs = 1) {
-        for (let epoch = 0; epoch < epochs; epoch++) {
-            let logs;
-            let peakMemory = 0;
-            await this.model.fitDataset(trainingData, {
-                epochs: 1,
-                validationData,
-                callbacks: {
-                    onBatchEnd: (_) => {
-                        const currentMemory = tf.memory().numBytes / 1024 / 1024 / 1024; // GB
-                        if (currentMemory > peakMemory) {
-                            peakMemory = currentMemory;
-                        }
-                    },
-                    onEpochEnd: (_, cur) => { logs = cur; }
-                },
-            });
-            if (logs === undefined) {
-                throw new Error("Epoch didn't gave any logs");
-            }
-            const { loss, acc, val_acc, val_loss } = logs;
-            if (loss === undefined || isNaN(loss) || acc === undefined || isNaN(acc)) {
-                throw new Error("Training loss is undefined or nan");
-            }
-            const structuredLogs = {
-                epoch,
-                peakMemory,
-                training: {
-                    loss: logs.loss,
-                    accuracy: logs.acc,
-                }
+    async *train(trainingData, validationData) {
+        const batches = await trainingData.iterator(); // tf.LazyIterator isn't an AsyncGenerator
+        let batchesLogs = List();
+        for (let batchNumber = 0; true; batchNumber++) {
+            const iteration = await batches.next();
+            if (iteration.done)
+                break;
+            const batch = iteration.value;
+            const batchLogs = {
+                batch: batchNumber,
+                ...(await this.#runBatch(batch)),
             };
-            if (validationData !== undefined) {
-                if (val_loss === undefined || isNaN(val_loss) ||
-                    val_acc === undefined || isNaN(val_acc)) {
-                    throw new Error("Invalid validation logs");
-                }
-                structuredLogs.validation = {
-                    accuracy: logs.val_acc,
-                    loss: logs.val_loss
-                };
-            }
-            yield structuredLogs;
+            tf.dispose(batch);
+            yield batchLogs;
+            batchesLogs = batchesLogs.push(batchLogs);
         }
+        const validation = validationData && (await this.#evaluate(validationData));
+        return new EpochLogs(batchesLogs, validation);
+    }
+    async #runBatch(batch) {
+        let logs;
+        await this.model.fitDataset(tf.data.array([batch]), {
+            epochs: 1,
+            verbose: 0, // don't pollute
+            callbacks: {
+                onEpochEnd: (_, cur) => {
+                    logs = cur;
+                },
+            },
+        });
+        if (logs === undefined)
+            throw new Error("batch didn't gave any logs");
+        const { loss, acc: accuracy } = logs;
+        if (loss === undefined || isNaN(loss))
+            throw new Error("training loss is undefined or NaN");
+        return {
+            accuracy,
+            loss,
+            memoryUsage: tf.memory().numBytes / 1024 / 1024 / 1024,
+        };
+    }
+    async #evaluate(dataset) {
+        const evaluation = await this.model.evaluateDataset(dataset.map((t) => {
+            switch (t) {
+                case null:
+                case undefined:
+                    throw new Error("nullish value in dataset");
+                default:
+                    return t;
+            }
+        }));
+        const metricToValue = Map(List(this.model.metricsNames).zip(Array.isArray(evaluation)
+            ? List(await Promise.all(evaluation.map((t) => t.data())))
+            : List.of(await evaluation.data()))).map((values) => {
+            if (values.length !== 1)
+                throw new Error("more than one metric value");
+            return values[0];
+        });
+        const [accuracy, loss] = [
+            metricToValue.get("acc"),
+            metricToValue.get("loss"),
+        ];
+        if (accuracy === undefined || loss === undefined)
+            throw new Error("some needed metrics are missing");
+        return { accuracy, loss };
     }
     predict(input) {
         const ret = this.model.predict(input);

package/dist/training/disco.d.ts

@@ -1,4 +1,4 @@
-import type { data, Logger, Memory, Task, TrainingInformation } from '../index.js';
+import { BatchLogs, data, EpochLogs, Logger, Memory, Task, TrainingInformation } from '../index.js';
 import { client as clients } from '../index.js';
 import type { Aggregator } from '../aggregator/index.js';
 import type { RoundLogs } from './trainer/trainer.js';
@@ -22,13 +22,25 @@ export declare class Disco {
     private readonly client;
     private readonly trainer;
     constructor(task: Task, options: DiscoOptions);
-    /**
-     * Starts a training instance for the Disco object's task on the provided data tuple.
-     * @param dataTuple The data tuple
-     */
-    fit(dataTuple: data.DataSplit): AsyncGenerator<RoundLogs & {
+    /** Train on dataset, yielding logs of every round. */
+    trainByRound(dataTuple: data.DataSplit): AsyncGenerator<RoundLogs & {
         participants: number;
     }>;
+    /** Train on dataset, yielding logs of every epoch. */
+    trainByEpoch(dataTuple: data.DataSplit): AsyncGenerator<EpochLogs>;
+    /** Train on dataset, yielding logs of every batch. */
+    trainByBatch(dataTuple: data.DataSplit): AsyncGenerator<BatchLogs>;
+    /** Run whole train on dataset. */
+    trainFully(dataTuple: data.DataSplit): Promise<void>;
+    /**
+    * Train on dataset, yield the nested steps.
+    *
+    * Don't forget to await the yielded generator otherwise nothing will progress.
+    * If you don't care about the whole process, use one of the other train methods.
+    **/
+    train(dataTuple: data.DataSplit): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs & {
+        participants: number;
+    }>>;
     /**
      * Stops the ongoing training instance without disconnecting the client.
      */

package/dist/training/disco.js

@@ -1,5 +1,8 @@
+import { List } from 'immutable';
+import { async_iterator } from '../index.js';
 import { client as clients, EmptyMemory, ConsoleLogger } from '../index.js';
 import { MeanAggregator } from '../aggregator/mean.js';
+import { enumerate, split } from '../utils/async_iterator.js';
 import { TrainerBuilder } from './trainer/trainer_builder.js';
 /**
  * Top-level class handling distributed training from a client's perspective. It is meant to be
@@ -58,35 +61,79 @@ export class Disco {
         const trainerBuilder = new TrainerBuilder(this.memory, this.task);
         this.trainer = trainerBuilder.build(this.client, options.scheme !== 'local');
     }
+    /** Train on dataset, yielding logs of every round. */
+    async *trainByRound(dataTuple) {
+        for await (const round of this.train(dataTuple)) {
+            const [roundGen, roundLogs] = async_iterator.split(round);
+            for await (const epoch of roundGen)
+                for await (const _ of epoch)
+                    ;
+            yield await roundLogs;
+        }
+    }
+    /** Train on dataset, yielding logs of every epoch. */
+    async *trainByEpoch(dataTuple) {
+        for await (const round of this.train(dataTuple)) {
+            for await (const epoch of round) {
+                const [epochGen, epochLogs] = async_iterator.split(epoch);
+                for await (const _ of epochGen)
+                    ;
+                yield await epochLogs;
+            }
+        }
+    }
+    /** Train on dataset, yielding logs of every batch. */
+    async *trainByBatch(dataTuple) {
+        for await (const round of this.train(dataTuple))
+            for await (const epoch of round)
+                yield* epoch;
+    }
+    /** Run whole train on dataset. */
+    async trainFully(dataTuple) {
+        for await (const round of this.train(dataTuple))
+            for await (const epoch of round)
+                for await (const _ of epoch)
+                    ;
+    }
     /**
-     * Starts a training instance for the Disco object's task on the provided data tuple.
-     * @param dataTuple The data tuple
-     */
+    * Train on dataset, yield the nested steps.
+    *
+    * Don't forget to await the yielded generator otherwise nothing will progress.
+    * If you don't care about the whole process, use one of the other train methods.
+    **/
     // TODO RoundLogs should contain number of participants but Trainer doesn't need client
-    async *fit(dataTuple) {
+    async *train(dataTuple) {
         this.logger.success("Training started.");
         const trainData = dataTuple.train.preprocess().batch();
         const validationData = dataTuple.validation?.preprocess().batch() ?? trainData;
         await this.client.connect();
         const trainer = await this.trainer;
-        for await (const roundLogs of trainer.fitModel(trainData.dataset, validationData.dataset)) {
-            let msg = `Round: ${roundLogs.round}\n`;
-            for (const epochLogs of roundLogs.epochs.values()) {
-                msg += `    Epoch: ${epochLogs.epoch}\n`;
-                msg += `        Training loss: ${epochLogs.training.loss}\n`;
-                if (epochLogs.training.accuracy !== undefined) {
-                    msg += `        Training accuracy: ${epochLogs.training.accuracy}\n`;
+        for await (const [round, epochs] of enumerate(trainer.fitModel(trainData.dataset, validationData.dataset))) {
+            yield async function* () {
+                let epochsLogs = List();
+                for await (const [epoch, batches] of enumerate(epochs)) {
+                    const [gen, returnedEpochLogs] = split(batches);
+                    yield gen;
+                    const epochLogs = await returnedEpochLogs;
+                    epochsLogs = epochsLogs.push(epochLogs);
+                    this.logger.success([
+                        `Round: ${round}`,
+                        `  Epoch: ${epoch}`,
+                        `    Training loss: ${epochLogs.training.loss}`,
+                        `    Training accuracy: ${epochLogs.training.accuracy}`,
+                        epochLogs.validation !== undefined
+                            ? `    Validation loss: ${epochLogs.validation.loss}`
+                            : "",
+                        epochLogs.validation !== undefined
+                            ? `    Validation accuracy: ${epochLogs.validation.accuracy}`
+                            : "",
+                    ].join("\n"));
                 }
-                if (epochLogs.validation !== undefined) {
-                    msg += `        Validation loss: ${epochLogs.validation.loss}\n`;
-                    msg += `        Validation accuracy: ${epochLogs.validation.accuracy}\n`;
-                }
-            }
-            this.logger.success(msg);
-            yield {
-                ...roundLogs,
-                participants: this.client.nodes.size + 1 // add ourself
-            };
+                return {
+                    epochs: epochsLogs,
+                    participants: this.client.nodes.size + 1, // add ourself
+                };
+            }.bind(this)();
         }
         this.logger.success("Training finished.");
     }

package/dist/training/trainer/trainer.d.ts

@@ -1,9 +1,8 @@
 import type tf from "@tensorflow/tfjs";
 import { List } from "immutable";
 import type { Model, Task } from "../../index.js";
-import { EpochLogs } from "../../models/model.js";
+import { BatchLogs, EpochLogs } from "../../models/index.js";
 export interface RoundLogs {
-    round: number;
     epochs: List<EpochLogs>;
 }
 /** Abstract class whose role is to train a model with a given dataset. This can be either done
@@ -29,5 +28,5 @@ export declare abstract class Trainer {
      * Start training the model with the given dataset
      * @param dataset
      */
-    fitModel(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<RoundLogs>;
+    fitModel(dataset: tf.data.Dataset<tf.TensorContainer>, valDataset: tf.data.Dataset<tf.TensorContainer>): AsyncGenerator<AsyncGenerator<AsyncGenerator<BatchLogs, EpochLogs>, RoundLogs>, void>;
 }

package/dist/training/trainer/trainer.js

@@ -1,4 +1,5 @@
 import { List } from "immutable";
+import * as async_iterator from "../../utils/async_iterator.js";
 /** Abstract class whose role is to train a model with a given dataset. This can be either done
  * locally (alone) or in a distributed way with collaborators.
  *
@@ -16,6 +17,8 @@ export class Trainer {
         this.model = model;
         this.#roundDuration = task.trainingInformation.roundDuration;
         this.#epochs = task.trainingInformation.epochs;
+        if (!Number.isInteger(this.#epochs / this.#roundDuration))
+            throw new Error(`round duration doesn't divide epochs`);
     }
     /**
      * Request stop training to be used from the Disco instance or any class that is taking care of the trainer.
@@ -28,25 +31,31 @@ export class Trainer {
      * @param dataset
      */
     async *fitModel(dataset, valDataset) {
-        if (this.training !== undefined) {
+        if (this.training !== undefined)
             throw new Error("training already running, cancel it before launching a new one");
+        try {
+            this.training = this.#runRounds(dataset, valDataset);
+            yield* this.training;
         }
-        await this.onRoundBegin(0);
-        this.training = this.model.train(dataset, valDataset, this.#epochs);
-        for await (const logs of this.training) {
-            // for now, round (sharing on network) == epoch (full pass over local data)
-            yield {
-                round: logs.epoch,
-                epochs: List.of(logs),
-            };
-            if (logs.epoch % this.#roundDuration === 0) {
-                const round = Math.trunc(logs.epoch / this.#roundDuration);
-                await this.onRoundEnd(round);
-                await this.onRoundBegin(round);
-            }
+        finally {
+            this.training = undefined;
         }
-        const round = Math.trunc(this.#epochs / this.#roundDuration);
-        await this.onRoundEnd(round);
-        this.training = undefined;
+    }
+    async *#runRounds(dataset, valDataset) {
+        const totalRound = Math.trunc(this.#epochs / this.#roundDuration);
+        for (let round = 0; round < totalRound; round++) {
+            await this.onRoundBegin(round);
+            yield this.#runRound(dataset, valDataset);
+            await this.onRoundEnd(round);
+        }
+    }
+    async *#runRound(dataset, valDataset) {
+        let epochsLogs = List();
+        for (let epoch = 0; epoch < this.#roundDuration; epoch++) {
+            const [gen, epochLogs] = async_iterator.split(this.model.train(dataset, valDataset));
+            yield gen;
+            epochsLogs = epochsLogs.push(await epochLogs);
+        }
+        return { epochs: epochsLogs };
     }
 }

package/dist/utils/async_iterator.d.ts

@@ -0,0 +1,11 @@
+import { List } from "immutable";
+/**
+ * Split yields from return value
+ *
+ * You need to consume the iterator to resolve the returned value
+ **/
+export declare function split<T, U>(iter: AsyncIterator<T, U>): [AsyncGenerator<T, U>, Promise<U>];
+/** Zip iterator with a infinite counter */
+export declare function enumerate<T, U>(iter: AsyncIterator<T, U> | Iterator<T, U>): AsyncGenerator<[number, T], U>;
+/** Run the whole iterator to get yielded & returned */
+export declare function gather<T, U>(iter: AsyncIterator<T, U>): Promise<[List<T>, U]>;

package/dist/utils/async_iterator.js

@@ -0,0 +1,63 @@
+import { List } from "immutable";
+// `Promise.withResolvers` not widely deployed
+function PromiseWithResolvers() {
+    let resolve, reject;
+    resolve = reject = () => {
+        // should not happen as Promise are run on creation
+        throw new Error("race condition triggered");
+    };
+    const promise = new Promise((res, rej) => {
+        resolve = res;
+        reject = rej;
+    });
+    return [promise, resolve, reject];
+}
+/**
+ * Split yields from return value
+ *
+ * You need to consume the iterator to resolve the returned value
+ **/
+export function split(iter) {
+    const [returnPromise, returnResolve, returnReject] = PromiseWithResolvers();
+    return [
+        (async function* () {
+            try {
+                while (true) {
+                    const v = await iter.next();
+                    if (!v.done) {
+                        yield v.value;
+                        continue;
+                    }
+                    returnResolve(v.value);
+                    return v.value;
+                }
+            }
+            catch (e) {
+                returnReject(e);
+                throw e;
+            }
+        })(),
+        returnPromise,
+    ];
+}
+/** Zip iterator with a infinite counter */
+export function enumerate(iter) {
+    return (async function* () {
+        for (let i = 0;; i++) {
+            const v = await iter.next();
+            if (v.done)
+                return v.value;
+            yield [i, v.value];
+        }
+    })();
+}
+/** Run the whole iterator to get yielded & returned */
+export async function gather(iter) {
+    let elems = List();
+    for (;;) {
+        const v = await iter.next();
+        if (v.done)
+            return [elems, v.value];
+        elems = elems.push(v.value);
+    }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@epfml/discojs",
-  "version": "2.1.2-p20240627125649.0",
+  "version": "2.2.0",
   "type": "module",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",