Repository navigation
gpu, model-hf, placement: read a checkpoint folder's parts in parallel within the host load budget - #787
Conversation
github-actions
Bot
added
gpu
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: MUST_FIX: parallel part loading adds host allocations that the placement memory check does not budget.
Must fix
1. Bound the owned parts retained before staging
crates/gpu/src/parts.rs:54 reads the entire run before staging, then keeps its parts in done until the run ends. These are owned allocations: Checkpoint::read:225 returns Part::made. On a discrete GPU with 1 GiB available after other host allocations and no swap, 256 normalized 32 MiB parts require 8 GiB of retained results [derived], before direct-read and packing buffers. The previous loop retained one part.
The load's HostNeed check:2193 runs before the part loader; its budget terms:783 do not reserve this run. Such a load passes the check and exhausts host memory instead of returning a load error.
Fix this before deployment: the failure loses the server process instead of returning a load error.
Restore the one-part lifecycle below until the common memory accounting reserves the full parallel read/normalization peak. Remove READ_AHEAD, read_ahead, its policy-only test, and the unused Planned/pool_map imports; update the load documentation. Verify loading under a constrained host-memory limit.
pub fn load<S: Source, E: std::fmt::Display>(
stream: &CudaStream,
source: &S,
rule: impl Fn(&str, GgmlType, &[u64]) -> Result<Take, E>,
card: impl Fn(&TensorInfo) -> Option<CardFormat>,
release: bool,
) -> Result<Weights, GpuError> {
let planned = source.plan(rule).map_err(|e| refused(e.to_string()))?;
let mut budget = 0usize;
for p in &planned {
budget = budget
.checked_add(part_format(&p.info)?.1)
.ok_or_else(|| refused("the parts' staging bytes pass usize".into()))?;
}
stream.context().bind_to_thread()?;
let mut staged = BTreeMap::new();
// Made after `staged` and dropped before it: its drop drains the copies, so no buffer of a
// refused load is freed while a copy still writes it.
let mut ring = UploadRing::new(stream, budget)?;
for p in &planned {
let part = source.read(p).map_err(|e| refused(e.to_string()))?;
let one = TensorPart {
info: part.info.clone(),
spans: part.spans(),
scale: part.scale.clone(),
};
let held = |name: &str| staged.contains_key(name);
let added = stage_parts(stream, &mut ring, &[one], held, &card)?;
staged.extend(added);
drop(part);
if release {
source.release(p).map_err(|e| refused(e.to_string()))?;
}
}
ring.finish(stream)?;
Ok(Weights::of_staged(staged))
}Non-blocking
⚠ Complete the final-head evidence: crates/model-hf/src/part.rs:273 changes checkpoint I/O. The body places the gate runs before the last two commits and omits the measured commit hashes and complete startup commands. Attach both run hashes, exact model paths, commands/output and selected gates on the final head so the startup claims can be reproduced.
Checked
Lifetime and budget: traced Folder38 through parts::load_placed, the owned checkpoint results, the pool's result collection, and the preceding HostNeed grant.
Fix: the replacement function matches the base loader; the upload ring copies source bytes into its own pinned slots before returning.
Tests: the read-ahead test checks the 8 GiB boundary, but has no host-headroom input.
Evidence: the PR body confirms the final-head gate omissions and lacks measured commit hashes and complete startup commands.
Scope: full review; no revert, no no-merge mirror, no ledger lessons. Escape catalog C1–C4 do not apply to the retained findings.
Notes
Rust/GPU suites were not executed during validation.
Card check: needs-gpu. This pull request changes crates/gpu-gates/src/bin/bloomery_serve.rs, crates/gpu-gates/src/bin/shared/serve_seats/mod.rs, crates/gpu-gates/src/bin/shared/serve_seats/qwen38_folder.rs … and carries no gate run on a card (CONTRIBUTING's evidence list); a maintainer adds gpu-verified after one.
🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)
host_filled's advice (madvise(MADV_HUGEPAGE) over the whole huge pages of a new allocation of 32 MiB or more) moves to bloomery-os as advise_huge_within, so the loader's own new buffers can take it too. The gpu block moves with it; the unsafe pins follow.
A checkpoint folder's read normalized each tensor from the shard's mapping, so every 4 KiB page of the 78 GB Qwen3.8 FP4-Local checkpoint was a fault on the reading thread, and the pages then sat in the page cache until the load dropped them again. Checkpoint::read now reads each stored tensor of a whole tensor with one direct read through engram's DirectFile into a buffer of its own, and normalizes from those bytes; a part of a tensor (a tensor-parallel rank's rows) still reads its rows through the mapping, and so does a shard whose mount takes no direct reads (DirectFile::open refuses it). The new buffers, and the normalized NVFP4 and FP8 blocks, are advised onto huge pages (os::advise_huge_within), so filling them faults once each 2 MiB.
append moved each NVFP4 group's 16 nibbles into GGUF's order one byte at a time through a per-byte iterator, and decoded every scale byte with powi. The group is now two 32-bit words whose nibbles are spread with three shift-and-mask steps, and the scale bytes' half values come from a table built once a tensor. A byte that is no UE4M3 scale is refused as before, naming the same row and group. The modelopt tests' reference values hold the bytes.
spread adds one expert at a time and asks the card's footprint after each one, and each ask replayed every upload's buffers through the allocator model. Qwen3.8 has 48 x 512 routed experts and about 3000 tensors, so one plan took about 3 s on GB10, and the seat plans five times before the load (the placement choice, the context rule, the host need and the final plan). Footprint keeps each step's buffers at its last rows, and what its small buffers did to the shared granules. A buffer of a granule or more takes whole granules wherever it falls in the order, so only the small buffers depend on the granules a step finds, and a step that finds them as it last did, with the same small buffers, leaves them as it last did. The footprint is the same bytes as the replay at any counts; the new test asks it at counts that grow, shrink and cross a granule.
parts::load read, normalized and staged one part at a time on one thread, and the host packings of the staged parts (the NVFP4 and MXFP4 repack, the E4M3 and BF16 MXFP8 planes) ran a row at a time on the same thread, so a 78 GB checkpoint load kept one core busy while the drive and the 48.9 GB/s copy lane waited. Qwen3.8 FP4-Local took about 270 s to load on GB10. The parts are now read in runs of at most 8 GiB, each run's parts at once over the resident worker pool (pool_map, each thread taking the next part as it finishes one), then staged in plan order, and the run's parts freed and their pages dropped over the pool. The host packings pack their rows in runs over the pool (pack_planes) into planes advised onto huge pages. A refusal is the first in plan or row order, with the same text as before. A source is Sync for its parts to be read at once.
69990e8 to
fb859cb
Compare
juliankang4
changed the title
There was a problem hiding this comment.
Since the last review:
- Must fix 1: fixed. crates/gpu/src/parts.rs:175 counts read/packing costs before allocation and checks each run's live room.
- Earlier hardware-evidence warning: still open at crates/model-hf/src/part.rs:272; details below.
Verdict: CLEAN: the revised admission checks address the retained-memory failure; no new must-fix finding survives.
Non-blocking
⚠ Complete final-head evidence: crates/model-hf/src/part.rs:272 changes checkpoint I/O. The body omits both measured commit hashes, exact model paths and complete startup commands/output. It states that the startup comparison and BF16 gate ran before the final changes. Attach the missing run details and selected gate results for the final head so the measurements can be reproduced.
Checked
Duplicates: diagnostics report zero dryer candidates in the searched crates; GPU and GPU-gates were excluded from that search.
Complexity: diagnostics flag seven functions, including new load_plan orchestration and added Checkpoint::read I/O branches.
Silent failure: read-cost arithmetic returns named overflow errors. Direct-read fallback accepts only invalid or unsupported operations; other I/O errors propagate.
Admission: plan.check precedes upload-ring allocation and source reads. Placed loads include read-ahead and staging in HostNeed; each run checks its own cost against live room.
Tests: assertions cover refusal before source.read, falling host room, per-run costs, batching limits and read-cost overflow. Removing the initial admission check reaches the test source's distinct error; using the whole-plan peak rejects the later run that the test expects to fit.
Gates: the unsafe-count change transfers one block from GPU to OS with a dated PIN and reason.
Escape catalog: C1–C4 do not match the retained warning.
Notes
Validation was static. Rust/CUDA suites and hardware gates were not executed during validation.
Card check: needs-gpu. This pull request changes crates/gpu-gates/src/bin/bloomery_serve.rs, crates/gpu-gates/src/bin/shared/serve_seats/minimaxm3_folder.rs, crates/gpu-gates/src/bin/shared/serve_seats/qwen38_folder.rs … and carries no gate run on a card (CONTRIBUTING's evidence list); a maintainer adds gpu-verified after one.
🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)
|
The review is done and the contribution check passed, so I'm merging this. |
Qwen3.8 모델을 GB10에서 서버로 시작할 때 걸리는 시간을 줄이고, 병렬 읽기가 쓰는 호스트 메모리를 로드 예산에 넣습니다.
What and why
Starting
bloomery-serveon a Qwen3.8 FP4-Local folder on a DGX Spark spends most of its time inparts::load. It read, normalized and staged one part at a time on one thread, the checkpoint reads were page faults on a memory-mapped shard, and the NVFP4 and MXFP4 repacks and the E4M3 and BF16 plane packing ran a row at a time. The expert rule in the planner also replayed every upload buffer of the card each time it added one expert.This PR removes those costs at their owners:
parts::loadreads a load's parts a run at a time over the worker pool and stages them in plan order. The repack and plane packing run on the pool as well (pack_planes). The first refused part or row is still the first in plan and row order, with the same message.Checkpoint::readinmodel-hfreads a whole stored tensor with the O_DIRECTDirectFileinto a huge-page buffer. A tensor-parallel part, and a mount that does not accept direct reads, still read through the mapping.nvfp4_group_to_block, and the scale table is built once per tensor.Footprintinplacementkeeps what no expert count changes, so the expert rule stops replaying all buffers per expert. A new test checks that its bytes equal the full replay at several counts.Source::read_bytesreports each source's owned bytes (direct input, normalization scratch, decoded scales, widened values, and the contiguous copy and scales of a GGUF rank cut).read_aheadcounts as anonymous memory, and the pinned upload ring counts as shadows. Admission of the whole plan happens before the ring is built or any part is read, and each run is then checked against its own cost and the room left at that moment, so a small later run is not refused because an earlier run peaked higher.gputoos::advise_huge_within, somodel-hfandgpushare it.model_fileowner that the Qwen and MiniMax seats both use.Release notes: Changed: Qwen3.8 FP4-Local folders start in 78.042 s instead of 225.325 s on a DGX Spark (median of three starts).
Code shape
DirectFile, the huge-page advice is one function inos, the NVFP4 group packing stays in its existing common owner, and the opened-folder cache is inmodel_filefor both seats.parts::load,Checkpoint::read,Footprint,model_file) and adds no per-model branch, flag or trait method.gpucopy of the huge-page advice is removed.DirectFile::opento the mapping happens only when the filesystem reports the direct read as invalid or unsupported. Every other I/O error and every short read is a named refusal, and a shard that is already mapped does not hide an open error. Cost sums and products are checked, and an overflow is a named error.cargo fmtandcargo clippyare clean on what I touched; the lint count did not go up. On the head before the last rebase,cargo clippy --all-targets -- -D warningson the touched crates failed only on an unusedNVFP4_OFFSETconstant incrates/qdot/src/lib.rs, which this diff does not touch and which main has since fixed. I did not re-run clippy on the final head. Thegpucrate cannot build on the Mac, so it was compiled on a node.Duplication found or removed (
path:line), or "none": none.What I ran
Startup time on a DGX Spark with
--host 127.0.0.1 --parallel 4 --ctx-size 16896,BLOOMERY_HOST_ROOM=96G, draft model off and greedy sampling. Startup runs from spawning the server to the end of a 32-token completion ofThe capital of France is. After each load, a 128-token greedy completion of the same long prompt was compared across arms. One warm load per arm was discarded, thenmain-1, head-1, main-2, head-2, main-3, head-3ran on the same node in one sitting, each after a 30 second quiet-CPU check. The model is the Qwen3.8 FP4-Local folder, and the "main" arm isupstream/mainat the time of measuring.Startup time fell by 65.365% in the same-node alternating sitting. The three starts were 230.423576 s, 224.398483 s and 225.324722 s on main and 78.036886 s, 78.041724 s and 79.302609 s on this change. The 128-token greedy text and the long-prompt guard output were equal across all arms.
Gate on the same DGX Spark:
Regression tests on an RTX 5090 machine. Each one was first run with its guard removed to confirm that it catches the removal, then with the guard restored:
cargo check --workspace --all-targetsandcargo check -p bloomery-gpu --all-targetspass.What I ran into
This does not block the PR.
gate-gpu-bf16run are on an earlier head of this branch. The head that goes up differs by a rebase onto the current main (the sharedFolderand packing owners), the per-run admission correction, and the restored format check ahead of the first read. A sufficient host room, as in this sitting, does not take the changed refusal branch. I did not re-measure the final head on the GB10, and I ran the regression tests and compile checks on it only on the RTX 5090 machine.parts::load.Hardware (if the change is for a machine)
One DGX Spark (GB10), driver 595.84 and CUDA 13.3 for the startup comparison and the gate, and one RTX 5090 machine for the regression tests.
AI help
AI 코딩 에이전트가 코드를 작성한 후 설명을 썼고, 테스트를 마쳤습니다. diff 리뷰와 PR 초안 리뷰는 제가 직접 했습니다.