Fit planner: deterministic per-file chunk budgets (device_memory_budget)

[gitbisector]

A fixed max_batch_bytes budget (#90) pays chunking cost on every shard even while device memory is still empty. The fit planner precomputes per-file budgets from the safetensors headers: whole-file loads while headroom is ample, budgets declining as resident bytes grow, chunking only the tail shards — and a plan-time BudgetInfeasibleError naming 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. The budget is an explicit integer, which makes it broadcast-safe: the same plan lands on every rank, with one extra depth unit for the in-flight receive tensor.

Choosing the budget is the caller's job, per review. The caller knows what else lives on the device and what every other rank has; the planner module documents the reserve heuristic and the all-reduce(MIN) recipe for multi-rank callers to lift:

free, _ = torch.cuda.mem_get_info(dev)
budget = free - max(free // 20, 1 << 30)   # reserve for allocator rounding + fixed pools
# multi-rank: all-reduce(MIN) budget before passing it, so every rank plans identically

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, 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 via transient_multiplier, which each copier reports for itself: CopierInterface.chunk_transient_multiplier defaults to raising NotImplementedError (same contract as set_chunk), UnifiedMemCopier mirrors its own reader-path selection, and NoGdsFileCopier returns 1. parallel_loader reaches it through the loader's registered copier class, so it depends on no specific copier module.

Measured: on a 46-shard / 160 GB checkpoint, the planner matched full-speed loading at every feasible budget.

[takeshi-yoshimura]

@gitbisector
Thanks for contributing this. I understand this is important especially for unified memory systems like DGX Spark to utilize more memory while avoiding OOM. Can you rebase recent changes and make this ready for review?

I would suggest that the auto budgeting could be implemented in the caller side rather than here because callers should have more knowledge of ranks' memory usages.

Also, parallel_loader.py should not depend on a specific copier module. Specifically, we need to remove the line from .copier.unified import chunk_transient_multiplier. CopierInterface should provide chunk_transient_multiplier and raise NotImplementedError in other copier backends like set_chunk.

[gitbisector]

Thanks — rebased onto current main (#90 having merged, this is now a single feature commit plus the review response), and both points addressed:

1. Caller-side budgeting. Agreed — the caller knows what else lives on the device and what the other ranks hold. device_memory_budget is now an int only: "auto", resolve_auto_budget, the FrameworkOpBase.get_mem_free op and its torch implementation, and the single-group guard that "auto" required are all gone (an explicit budget is already identical on every rank, so the plan stays deterministic without it). The reserve heuristic and the all-reduce(MIN) recipe moved to the _planner module docstring for callers to lift. frameworks/ is back to byte-identical with main.

2. No copier-specific import in parallel_loader. chunk_transient_multiplier is now on CopierInterface with a default that raises NotImplementedError, the same contract as set_chunk — a copier that stages a chunk more than once has to say so, or the plan under-counts and OOMs. UnifiedMemCopier keeps mirroring its own reader-path selection (1× on O_DIRECT, 2× on the mmap+pin fallback); NoGdsFileCopier returns 1, since its bounce-buffer pool is fixed-size and only the chunk buffer scales. The planner reaches it through the loader's copier class, registered alongside each factory (register_copier_constructor(type, cls) / get_copier_class), so parallel_loader imports no copier module at all — and gds/3fs now fail at plan time with a clear message instead of silently borrowing the unified copier's number.

Verified on a DGX Spark (GB10, CUDA 13.1, torch 2.10): tests/unit is 184 passed on this branch vs 164 on main, plus an end-to-end check on cuda:0 confirming a device_memory_budget load is byte-identical to a plain load, the multiplier resolves per copier and per reader path, and an infeasible budget is rejected at plan time naming the shard.

(Unrelated: tests/unit/threefs segfaults in the mock reader for me, but it reproduces identically on unmodified main — happy to open a separate issue if it isn't already known.)

[takeshi-yoshimura]

Merged. thanks! I have quickly opened the issue #99, but I believe that is low priority as of now.