feat: overlap chunk loading, gpu transfer and training #131
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
Introduces a new chunk loader that overlaps disk I/O, CPU->GPU transfer, and training using a 3-thread pipeline architecture: - cpu_loader thread: disk -> cpu_queue (pinned memory) - gpu_transfer thread: cpu_queue -> gpu_queue (device tensors) - main thread: gpu_queue -> training Key changes: - Add PipelineChunkLoader class with configurable gpu_prefetch parameter - Add comprehensive test suite including 10s timing stress test - Update load_dataset() to use PipelineChunkLoader - Enable double buffering (gpu_prefetch=2) in latent_stag.py Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
- add PipelineProfiler class for recording pipeline events - profile disk_load, gpu_transfer, gpu_queue_wait, and epoch events - save trace to pipeline_trace.json for viewing in chrome://tracing - profiler is disabled by default, enabled in latent_stag.py Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
adds "train" event on main thread to show when training is happening, making it easy to visualize overlap with data loading/transfer. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
wraps each chunk's processing (get_next_chunk + train) with a "chunk" event, making it easy to see chunk boundaries in the trace. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
profiles the first 5 epochs to capture warmup and steady-state behavior while keeping file size manageable (~32KB vs ~650KB for 100 epochs). Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
set gpu_prefetch=2 to overlap cpu->gpu transfer with training. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
deletes: - chunk_loader.py (replaced by pipeline_chunk_loader.py) - chunk_loader_test.py - test_prefetch_youtube.py - CHUNKED_STREAMING_INTEGRATION.md (obsolete) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Collaborator
Author
|
This doesn't make a huge difference, but we have full overlap now of loading from disk, transfer to gpu mem, training. on the a100 machines:
Training time dominates. |
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
RandomChunkLoaderwithPipelineChunkLoaderusing 3-stage parallelismgpu_prefetch=2)Architecture
cpu_loader thread: disk -> cpu (pinned memory) -> cpu_queue
gpu_transfer thread: cpu_queue -> gpu -> gpu_queue
main thread: gpu_queue -> training
Profiling
pipeline_trace.json)epoch,chunk,train,gpu_queue_wait,disk_load,gpu_transferchrome://tracingor https://ui.perfetto.dev/