Sub-file chunk loading bounded by max_batch_bytes#90
Conversation
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>
efde7b1 to
00ceb91
Compare
|
@gitbisector 1. APIsAdding max_batch_bytes as a public API makes sense, but I would prefer not to expand the existing public surface beyond it without a strong need. Since SafeTensorsMetadata is exported, could we move plan_chunks to an internal planner, and preserve the current get_tensors signature via a private helper such as _get_tensors? Please also add set_chunk to CopierInterface with a default implementation that raises NotImplementedError since every copier should have the extension as separate PRs. set_chunk_plan also appears to be internal orchestration and could be renamed _set_chunk_plan. 2. AI-discovered issuesAn AI-assisted test found that stopping iteration early can leave the producer thread blocked when chunking expands a shard into multiple batches. This can be reproduced with the tiny test fixture, whose largest tensor is 256 bytes: it = ParallelLoader(
...,
queue_size=-1,
max_batch_bytes=256,
).iterate_weights()
next(it)
it.close()
In ext.cpp, the new O_DIRECT worker threads call CUDA runtime APIs without first selecting the loader’s target device. Since the current CUDA device is thread-local and defaults to device 0, loaders targeting another device may still create contexts and associate pinned allocations with GPU 0. Could we pass the device ID to |
Stopping iterate_weights() mid-shard (generator close or an exception in
the consumer) left the non-daemon producer thread parked forever: with
queue_size <= 0 it blocks in consumer_processed.wait(), and with a full
batch_queue it blocks in put(); neither observed stop_event. Reproducible
once max_batch_bytes expands a shard into several chunk-batches:
it = ParallelLoader(..., queue_size=-1, max_batch_bytes=256).iterate_weights()
next(it)
it.close() # join times out; stranded thread prevents process exit
Make both blocking points cancellation-aware (bounded waits re-checking
stop_event), and have cleanup wake consumer_processed and drain the queue
so the producer exits promptly. Batches nobody will consume have their
file buffers closed -- by the producer on a failed put, by cleanup for
items left in the queue -- so device memory is released.
Regression test closes after the first tensor and asserts the producer
thread terminates.
Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: git bisector <gitbisector@gmail.com>
The O_DIRECT worker threads called CUDA runtime APIs without selecting a device; the current device is thread-local and defaults to 0 in a fresh thread, so loaders targeting another GPU created contexts and issued copies against device 0. Pass the loader's device index down from the unified copier (as nogds_file_reader already does) and cudaSetDevice it at the start of each worker; device_id < 0 (cpu device) leaves the default untouched. Allocate the shared pinned-buffer pool with the portable flag (cudaHostAllocPortable / hipHostMallocPortable) so a buffer first pinned under one device's context remains valid pinned memory when a later call targets a different device. The trailing cudaDeviceSynchronize moves into the workers, whose current device is the target -- on the calling thread it could synchronize the wrong device. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: git bisector <gitbisector@gmail.com>
Review follow-up: SafeTensorsMetadata is exported, so its surface stays as it was -- plan_chunks moves to the internal _planner module (it is a loading strategy, not metadata behavior) and get_tensors keeps its original signature, with the names-restricted variant as _get_tensors for copiers instantiating compact chunk buffers. The loader's set_chunk_plan becomes _set_chunk_plan: producer-side orchestration, not for callers. set_chunk joins CopierInterface with a default that raises NotImplementedError, replacing the loader's hasattr probe: unlike set_byte_ranges (whose no-op default is correct), a silently ignored chunk plan would break the memory bound and multiply full-file reads, so refusal is the only safe default. Partial-read copiers (nogds, unified) override it. ParallelLoader(max_batch_bytes=...) remains the one public entry point. Also rewords the early-close join log: the producer exits once any in-flight copy completes, so its survival past the join timeout is not an error in itself. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: git bisector <gitbisector@gmail.com>
|
Thanks for the review — all three addressed:
|
The thread-accounting assertion compared against all new threads, so tqdm's TMonitor daemon singleton -- which first spawns inside the test when the file runs as its own pytest process, as CI does -- was reported as a stranded producer. Count only non-daemon threads: the stranded producer is non-daemon, which is the failure mode under test. Also close the loader before the assertion so a failure cannot hold the loader (and its 16MB bounce buffer) alive through pytest's traceback refs and trip the bounce_buffer_bytes check of a later test, which is how this surfaced as a second CI failure. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: git bisector <gitbisector@gmail.com>
2 participants
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.