Draft: GPT2 training on MCQ med data - #1111
Conversation
mina5rovic
changed the title
fd4b676 to
3014cf5
Compare
…k for model to contxt 512
JulienVig
left a comment
JulienVig
left a comment
There was a problem hiding this comment.
1st part of my review, I'll finish at a different time, feel free to start addressing these comments
| id: NodeID; | ||
| waitForMoreParticipants: boolean; | ||
| payload: serialization.Encoded; | ||
| payload?: serialization.Encoded | null; |
There was a problem hiding this comment.
Because payload is null only at round 0, when no global weights exist yet. The server sets it to undefined, but MessagePack decodes undefined as null, so the interface must allow null. It was one of the silent error, took some time to catch.
There was a problem hiding this comment.
Shouldn't the first payload sent be the base model that the clients use to train on their local data? I saw that your PR changed this to send null on the first round but I don't see why
There was a problem hiding this comment.
The client already downloads model via super.connect() before opening the federated WebSocket (before round 0). On develop, the same weights are then sent again in NewFederatedNodeInfo, so my change avoids that duplicate transfer, which is costly for GPT-2. I can put it back but I'm not sure it is optimal.
| const model = this.model as unknown as IterationTrainableTextModel; | ||
| if (typeof model.trainNextBatches !== "function") | ||
| throw new Error("model does not support iteration-based training"); | ||
|
|
There was a problem hiding this comment.
I find this part really inconvenient because it's not generic, it only addresses on particular case. Here's an idea to work around this, let me know what you think
- Implement a new dataset method that splits the dataset after a specified number of steps and returns 1) one dataset with only the next specified number of batches and 2) the remainder of the dataset after these batches. e.g., const [nextBatches, datasetRemainder] = dataset.subset(maxBatchCount)
- From runIterationRounds, I think this allows you to do reuse the existing #runRound, e.g. this.#runRound(nextBatches, roundValidationDataset)
- With this method I don't think you need to deal with all the iterators
Let me know what you think and whether this is feasible
There was a problem hiding this comment.
I agree that the GPT-specific cast and trainNextBatches are not ideal, and moving this logic out of the model would make it more generic.
My only concern with dataset.subset(maxBatchCount) is that Dataset is currently immutable and can be iterated multiple times, while returning a remainder would make it stateful. Using take and drop could also become expensive because we would keep iterating over the already consumed batches.
Maybe we can keep the iterator in the trainer, wrap the next batches in a temporary dataset, and pass it to the existing model.train(...). This would remove the GPT-specific method without changing the current Dataset behavior. We would still need to be careful when reusing #runRoundByEpoch, since it trains for roundDuration full passes while the iteration-based version should consume each batch only once.
Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com>
Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com>
Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com>
Co-authored-by: Julien Vignoud <33122365+JulienVig@users.noreply.github.com>
2 participants
No description provided.