Fit planner: deterministic per-file chunk budgets (device_memory_budget)#91
Draft
gitbisector wants to merge 4 commits into
Draft
Fit planner: deterministic per-file chunk budgets (device_memory_budget)#91gitbisector wants to merge 4 commits into
gitbisector wants to merge 4 commits into
Conversation
gitbisector
force-pushed
the
pr-b-fit-planner
branch
from
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>
…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. - planner.py (new, pure): 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
7004e7a to
60984d9
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.