performance(weight-free): link deferred weights from Hugging Face checkpoints#643
Draft
rebel-jongho wants to merge 2 commits into
Draft
performance(weight-free): link deferred weights from Hugging Face checkpoints#643rebel-jongho wants to merge 2 commits into
rebel-jongho wants to merge 2 commits into
Conversation
…e license header - test_weight_source: remove the two overlapping tests (generated-state save round-trip already covered by the artifact integration test; base-config storage already covered by the serialization test) and prune now-unused imports. - Trim over-length comments in modeling.py and the run_qwen3 example. - Replace the truncated license header in weight_source.py with the full Apache 2.0 header used across the package. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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.
Pull Request Description
Type of Change
Changes Overview
This wires weight-free compilation into optimum-rbln. Exporting with
weight_free=Trueexcludes the state dict from the compiled artifact and streams the weights in from the original Hugging Face checkpoint at load time.weight_sourceutility (src/optimum/rbln/utils/weight_source.py): maps the compiler's weight names to the original checkpoint keys via tensor identity (data_ptr, storage offset, numel, dtype, shape, stride). Tensors that exist in the HF checkpoint are streamed from safetensors per pool-window; only tensors generated or derived during compilation are persisted separately torbln_generated_tensors.safetensorsnext to the artifact.modeling_decoderonly.py): at load, each runtime's required windows come fromruntime.weight_load_plan(), andruntime.load_weight_window(chunk, graph_ids)fills only those windows.weight_sourcemetadata (model_id, revision, subfolder, variant) drives asnapshot_download, so weights can be re-fetched from the original checkpoint independently of the export session.configuration_utils.py): addsweight_free,weight_source,weight_name_map, andgenerated_weight_maptoRBLNModelConfig.modeling_base.py,modeling_decoderonly.py): the base class carries_supports_weight_free=Falseplus export/load hooks; decoder-only opts in with_supports_weight_free=True. Requestingweight_freeon an unsupported model is rejected clearly, and quantization/LoRA combinations are guarded as not-yet-supported.examples/text-generation/run_qwen3.py): exercises the weight-free path end to end.Motivation and Context
A compiled artifact is dominated by its weights, which makes distribution and transfer expensive. Weight-free separates the weights from the artifact and injects them from the original checkpoint at load time, shrinking the artifact dramatically. For example, Qwen3-0.6B is a 2.47 GB baked artifact but only a 33.8 MB weight-free artifact plus a 264 KB generated-tensor file.
This feature depends on the runtime APIs (
weight_load_plan/load_weight_window) and the state-dict-excluding compile path introduced by the rebel-compiler weight-key PR (RBLN-SW/rebel_compiler#11982), so it should be validated and merged together with that compiler PR.Validation
tests/test_weight_source.py(new, 212 lines): unit-covers identity mapping, generated-tensor persistence, window streaming, and the fresh-load path.Related Issues
Conventional commit