perf: stream training data from CPU to GPU per batch #118
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
validation metrics are now only computed via diagnostics (rollout stats). this removes the redundant per-epoch validation train_step call. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- keep train data on CPU with pin_memory for fast async transfer - add extract_batch_windows() to extract and transfer batch windows - modify make_batches_random() to yield windowed data on GPU - update train_step functions to use windowed data directly - validation data stays on GPU (small enough to fit) this reduces GPU memory usage by not loading entire dataset to GPU. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
the ems warmup feature was not being used and added complexity. removed the config field, validation logic, and training loop handling. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
…ochs warmup epochs run before the main training loop and are additional to the configured 'epochs' count. removed misleading validation that required warmup < epochs. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
replaces synchronous CPU->GPU transfer with DataLoader that prefetches batches in parallel worker processes. this hides transfer latency by preparing next batches while GPU trains on current batch. - add WindowDataset class for contiguous slice access (faster than fancy indexing) - add create_dataloader() with RandomSampler for infinite random sampling - use persistent_workers and pin_memory for optimal throughput - separate augmentation index generation (runs on GPU after batch arrives) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Fixes CUDA multiprocessing issue where fork start method causes silent crashes during DataLoader worker initialization. Spawn is safer with CUDA tensors. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Moves data and stim tensors to shared memory before creating the dataset. This allows spawned worker processes to access the data without copying, fixing bus errors and semaphore leaks. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Collaborator
Author
|
This is just too slow. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
ems_warmup_epochsfeatureDetails
Training data and stimulus are kept on CPU with
pin_memoryenabled for fast async transfer. Each batch extracts a window and transfers it to GPU. Validation data stays on GPU since it's small.This reduces GPU memory usage by not loading the entire training dataset to GPU memory upfront.
Test plan
🤖 Generated with Claude Code