fix: convert_olmoe_merged.py no longer needs the whole checkpoint in RAM#543
Merged
Merged
Conversation
The old --repo path called snapshot_download (pulls the whole checkpoint to local disk first) then loaded every shard fully into RAM via load_file before processing -- for OLMoE-1B-7B's ~14GB bf16 checkpoint, that's a real problem on a RAM-constrained machine (this was hit directly: 21GB total RAM, most already in use by other processes). Rewrites both paths to be lazy: - --model (local dir): opens shards with safetensors.safe_open and reads tensors on demand instead of safetensors.torch.load_file's eager full-shard load. - --repo: downloads and processes ONE source shard at a time via hf_hub_download, deleting it immediately after extraction -- peak extra disk usage is one shard, not the full checkpoint. An expert's three projections can land in different shards; incomplete triples are held in a small in-RAM dict until their last projection arrives, then merged, quantized, and flushed. - Resumable: existing output shards are scanned (header-only, via safe_open) and already-converted tensors are skipped on a rerun, so an interrupted conversion picks up where it left off instead of restarting. Same output format and CLI as before (--repo/--model/--out), no behavior change for anyone already using it successfully -- this only changes how much RAM/disk it needs while doing so. Verified end-to-end: converted allenai/OLMoE-1B-7B-0125-Instruct on a 21GB-RAM machine (previously not attempted with the old eager-load version), all 1024 experts (16 layers x 64) present and complete, olmoe.c loads and runs correctly against the result.
opxyc
changed the title
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
2 participants
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.
Summary
convert_olmoe_merged.py --repocalledsnapshot_download(pulls the whole checkpoint to disk first) then loaded every shard fully into RAM viaload_filebefore processing. For OLMoE-1B-7B's ~14GB bf16 checkpoint that's a real problem on a RAM-constrained machine — hit directly on a 21GB-RAM box already running other things. Rewrites both--modeland--repopaths to be lazy (safetensors.safe_open, on-demand reads) and, for--repo, to download/process/delete one source shard at a time — peak extra disk is one shard, not the full checkpoint. Also resumable: existing output tensors are detected (header-only scan) and skipped on a rerun. Same CLI and output format as before; no behavior change for anyone already using it successfully, only less RAM/disk needed while doing so.Validation
make -C c check— same disk-quota caveat as the sibling PRs (pre-existing, unrelated test, not caused by this change); everything else passes.make -C c cuda-test(if applicable) — N/A.allenai/OLMoE-1B-7B-0125-Instructon a 21GB-RAM machine (not previously possible with the eager-load version), all 1024 experts (16 layers × 64) present and complete,olmoe.cloads and runs correctly against the result.Compatibility