Sub-file chunk loading bounded by max_batch_bytes#90
Draft
gitbisector wants to merge 3 commits into
Draft
Conversation
This was referenced Jul 7, 2026
gitbisector
force-pushed
the
pr-a-max-batch-bytes
branch
from
dee636d to
efde7b1
Compare
- copier/{nogds,unified}: set_chunk() allocates only the chunk span and
materializes just the chunk's names (compact buffer; offset remap reuses
the existing copy_start_offset arithmetic); the set_byte_ranges path is
unchanged.
- common.get_tensors: optional names subset (required for compact buffers).
- common.SafeTensorsMetadata.plan_chunks(): byte-budget partitioner; a tensor
is the atomic load unit so the budget must cover the largest kept tensor.
Chunks emit the kept tensors' gap-merged runs (not one coalesced span), so
a tensor_filter'd chunked load reads only kept bytes -- the compact buffer
still covers the chunk span, but I/O drops to kept bytes (GB10, V4-Flash
EP-slice, uniform 2GB budget: 15.4s -> 11.6s).
- config.LoaderConfig.max_batch_bytes knob.
Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: git bisector <gitbisector@gmail.com>
…fers - Port dma_load_runs (multithreaded O_DIRECT range reader) into ext.cpp + pybind (bypasses page cache, drives NVMe queue depth; the single-thread pin path is page-cache-bound ~2.5 GB/s and does NOT scale with threads -- measured). FASTSAFETENSORS_DMA_THREADS knob (default 8). - unified.submit_io: prefer dma_load_runs(base_off, starts, ends, nthreads) for both full and compact-chunk buffers; fall back to mmap+pin_memory if unavailable. - Gate the fast path off network filesystems: O_DIRECT forfeits kernel readahead / client caching there, where buffered mmap+pin performs better. get_fs_type() (longest-prefix /proc/mounts match) decides; log once per fs type; FASTSAFETENSORS_ODIRECT=1/0 forces either way and DMA_THREADS=0 disables the reader entirely. Validated GB10, DeepSeek-V4-Flash shard: byte-identical across budgets; peak buffer tracks max_batch_bytes (0.25GB->8%, 1GB->30%); O_DIRECT lifts full 2.26->4.7 GB/s, chunked up to 7.1 GB/s. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: git bisector <gitbisector@gmail.com>
_create_batches expands each file-batch into aligned chunk-batches (chunk j = each rank's j-th chunk of its shard, None once a rank runs out) so every rank issues the same broadcast sequence in lockstep. Each shard stays owned by one rank and loads in chunks over successive batches; _load_single_batch sets the per-file chunk plan (set_chunk_plan) so copy_files_to_device uses compact O_DIRECT reads. max_batch_bytes threaded through ParallelLoader + PipelineParallel. Chunk plans are refused (NotImplementedError naming the copier) on copiers without set_chunk (gds, dstorage): silently allocating the full data section per chunk-batch would break the memory bound AND multiply full-file reads. Reusable pinned-buffer pool for the chunked reader: many small dma_load_runs calls made per-chunk cudaHostAlloc/cudaFreeHost dominate; recycle 16MB pinned bounce buffers via a mutex-guarded free-list shared across calls (GB10, V4-Flash EP-slice rank0: 1GB budget 26.1 -> 16.6s; 2GB 18.1 -> 15.4s; baseline 12.1 -> 10.2s). Batch-spec -> (rank_file_map, chunk_plan) mapping factored into PipelineParallel._spec_to_maps. Validated: full DeepSeek-V4-Flash TP=2 across 2x GB10 via ParallelLoader. baseline (off): 69187 tensors, peak 7.21 GB/rank max_batch_bytes=1GB: 69187 tensors, peak 2.22 GB/rank (-69%) overlapping sampled tensors byte-identical (64/64, compared by name). Related: foundation-model-stack#71 Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: git bisector <gitbisector@gmail.com>
gitbisector
force-pushed
the
pr-a-max-batch-bytes
branch
from
efde7b1 to
00ceb91
Compare
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
Implements the sub-file batching discussed in #71 (design sketch by @ABNER-1): decouple peak device memory from shard file size by loading each file in byte-budgeted chunks instead of whole shards. Three commits: the chunk-planning mechanism, a multithreaded O_DIRECT chunk reader, and the ParallelLoader integration.
Mechanism
plan_chunks()partitions each file's tensors into chunks undermax_batch_bytes, with the largest single tensor as the atomic floor. Chunks carry the kept tensors' gap-merged runs, so atensor_filter'd load reads only kept bytes (composes with the byte-range selection from Add caller-side byte-range selection (sub-file reads) + expert-parallel filter #81).set_chunk); copiers without it (gds, dstorage) refuse chunk plans loudly rather than silently breaking the bound.max_batch_bytes=None(default) is behavior-identical to today.dma_load_runs, default 8 threads) reads runs straight into the compact buffer, with a reusable 16 MB pinned-buffer pool; it's gated off network filesystems where buffered mmap+pin wins, with env overrides.Measurements (DGX Spark GB10, DeepSeek-V4-Flash): peak transient buffer tracks the budget (1 GB budget → 30% of whole-file peak; full TP=2 load: 7.21 → 2.22 GB/rank, −69%); O_DIRECT lifts read throughput 2.26 → up to 7.1 GB/s; byte-identical results across budgets.
Opened as a draft per the #71 discussion, to settle API shape against real code. Follow-up PR (stacked): a deterministic fit planner that removes the fixed-budget chunking cost where headroom makes it unnecessary.