Streaming Subgraph Prefetching - Overlap Disk I/O with Computation

[kylesayrs]

Objective

Overlap disk I/O for the next subgraph with computation on the current subgraph to reduce overall runtime for large model processing.

Current Behavior

Processing currently follows a strictly serial pattern:

read weights
transfer weights
compute
read next weights
transfer next weights
compute
...

For sufficiently large models, disk reads can contribute significantly to total runtime.

Proposed Improvement

Since the identity of the next subgraph is already known, its parameters should be prefetched into CPU memory while the current subgraph is being processed on the GPU:

GPU:   [ compute layer N       ][ compute layer N+1     ]
CPU:           [ transfer N+1 ]       [ transfer N+2 ]
Disk:  [ read N+1           ][ read N+2           ]

Proposed Implementation

prefetcher = Prefetcher(...)

prefetcher.submit(subgraphs[0])

for i, subgraph in enumerate(subgraphs):
    prefetched = prefetcher.wait(subgraph)

    if i + 1 < len(subgraphs):
        prefetcher.submit(subgraphs[i + 1])

    move_to_device(prefetched, main_device)
    subgraph(batch)

Key Considerations

• Prefetching should materialize parameters on CPU rather than triggering their normal onload path to the GPU
• Prefetching should be bounded to avoid excessive CPU memory usage
  • Initially: current on GPU, next on CPU, rest on disk
• Should not regress performance on systems where storage bandwidth is not the bottleneck
• Requires a mechanism to explicitly materialize offloaded parameters onto CPU while keeping them associated with modules

Expected Benefits

• Hide disk-read latency behind GPU computation
• Reduced total processing time for very large models
• Improved utilization of system resources by enabling concurrent operations across GPU, CPU, and disk

[aayush7511]

Hi, can I work on this?