Fit planner: deterministic per-file chunk budgets (device_memory_budget)#91
Draft
gitbisector wants to merge 8 commits into
Draft
Fit planner: deterministic per-file chunk budgets (device_memory_budget)#91gitbisector wants to merge 8 commits into
gitbisector wants to merge 8 commits into
Conversation
gitbisector
force-pushed
the
pr-b-fit-planner
branch
from
July 7, 2026 02:33
aa7f213 to
7004e7a
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-b-fit-planner
branch
from
July 7, 2026 02:44
7004e7a to
60984d9
Compare
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>
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>
…budgets device_memory_budget (int bytes | "auto") bounds resident tensors + transient buffers at peak, per rank. A static fit plan -- precomputed from safetensors headers plus one free-memory query -- gives every file the largest budget the bound allows: whole-file loads while headroom is ample (the existing fast path, zero chunking overhead), per-file budgets declining as cumulative resident bytes grow, chunking only where the fit requires it, and a plan-time BudgetInfeasibleError naming the file and required bytes when the model cannot fit. Deterministic: same files, filter, budget -> same plan. - fit planner (pure, folded into the internal _planner module): FileWeightStats, pipeline_depth, collect_file_stats, plan_file_budgets, resolve_auto_budget (reserve = max(5% free, 1 GiB)); bound lemma in the module docstring. - accumulate_resident flag: True = consumer keeps yielded tensors; False = destinations preallocated -> uniform budget/depth. - frameworks: get_mem_free(dev) op (torch: cuda.mem_get_info / sysconf). - parallel_loader: per-file budgets feed the existing chunk-batch machinery; header reads done once and reused for planning and chunk expansion. - config: device_memory_budget field; forward both kwargs regardless of use_pipeline. Explicit integer budgets work with broadcast loading: the same plan lands on every rank (header reads are deterministic); one extra unit of pipeline depth accounts for the in-flight broadcast receive tensor. "auto" stays single-group only: per-rank free-memory readings diverge and would deadlock the lockstep broadcast -- callers owning a process group should all-reduce(MIN) free memory and pass the result. Transient cost is path-dependent (measured on GB10 unified memory, 1-2 GiB chunks): the O_DIRECT reader costs ~1x span per live chunk plus a fixed ~150 MB thread pool (absorbed by the auto reserve), but the mmap+pin_memory fallback additionally pins the chunk's pages until wait_io -- on unified memory both draws share one physical pool, so each live chunk costs ~2x span. plan_file_budgets takes transient_multiplier; chunk_transient_multiplier() in the unified copier mirrors submit_io's path selection and the parallel loader wires it in. - tests: planner math, infeasibility, auto reserve, multiplier (halving, 2x infeasibility, validation), randomized replay asserting peak <= budget, and a CPU end-to-end budgeted load byte-identical to a plain load. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: git bisector <gitbisector@gmail.com>
gitbisector
force-pushed
the
pr-b-fit-planner
branch
from
July 24, 2026 22:26
60984d9 to
05b9284
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.
Stacked on #90 (
max_batch_bytessub-file chunk loading). A fixed budget pays chunking cost on every shard even while device memory is still empty. The fit planner precomputes per-file budgets from the safetensors headers plus a single free-memory query at plan time: whole-file loads while headroom is ample, budgets declining as resident bytes grow, chunking only the tail shards — and a plan-timeBudgetInfeasibleErrornaming the file and required bytes, instead of an OOM mid-load.Deliberately deterministic — same files, filter, and budget produce the same plan; no runtime feedback — staying out of the adaptive-tuning territory flagged in #71. Explicit integer budgets are broadcast-safe (identical plan on every rank, one extra depth unit for the in-flight receive tensor);
"auto"is single-group only, with the all-reduce(MIN) recipe documented for multi-rank callers.Transient cost is path-dependent (measured on GB10 unified memory, 1–2 GiB chunks): the O_DIRECT reader costs ~1× span per live chunk plus a fixed ~150 MB thread pool (absorbed by the auto reserve), but the mmap+pin fallback also pins the chunk's pages until the copy completes — on unified memory both draw from one physical pool, so those chunks cost ~2× span. The planner charges accordingly:
transient_multiplier, chosen automatically by mirroring the copier's own path selection.Measured: on a 46-shard / 160 GB checkpoint, the planner matched full-speed loading at every feasible budget.