Enable weightfree for causallm1 - #1247
Merged
ochougul merged 15 commits intoAug 7, 2026
Merged
Conversation
…o path) Adds weight-free export support to QEfficient: export ONNX graph structure without embedding weights, then let the QAIC compiler load weights directly from the original safetensors checkpoint at compile time. ## New package: QEfficient/exporter/weight_free/ - core.py: export_weight_free_onnx(), _build_meta_qeff_model(), _find_checkpoint_key() (6-strategy lookup including fused-expert splits and gate/router renames), _promote_initializers_and_build_spec(), load_weight_free_ort_inputs() - spec.py: WeightSpec, WeightSpecInput, WeightSpecLocation, save/load helpers - transforms.py: CheckpointTransformPipeline, DtypeConversionCheckpointTransform, MoEExpertStackingCheckpointTransform, MoEFusedExpertSplitCheckpointTransform, GptOssMxfp4ExpertDequantSplitCheckpointTransform - examples/text_generation/weight_free/: end-to-end example scripts ## Core wiring (modeling_qeff.py, modeling_auto.py) - QEFFBaseModel: _checkpoint_transforms class attribute, weight_spec_path instance attribute, _export_via_weightfree() method - _export(): use_weight_free_export param; guards _model_offloaded_check(); dispatches to weight-free path; saves weight_spec.json alongside ONNX - get_onnx_path() / _compile(): thread use_weight_free_export through chain; pop use_weight_free_export from compiler_options so it never reaches qaic-compile - QEFFAutoModelForCausalLM: _checkpoint_transforms pipeline, use_weight_free_export on export() + compile(); same bs/kv_cache_shape adjustments as dynamo=True ## Infrastructure fixes - torch_patches.py: add temporarily_disable_nested_compile_regions() — used by weight-free core.py for flat dynamo graph path - rms_norm.py: GemmaCustomRMSNormAIC forward-time weight+1.0 (meta-device safe; replaces __qeff_init__ copy_ which fails on meta tensors) - sampler_utils.py: dynamic_shapes as optional 4th return value - cache_utils.py: add device=position_ids.device / device=kv_position_ids.device to all torch.arange calls — required for meta-device tracing - modeling_attn_mask_utils.py: same device= fix for torch.arange in _create_causal_mask ## Model fixes: cos/sin device alignment for meta-device tracing llama, falcon, mistral, gemma, gemma2, granite, olmo2, qwen2, qwen3, qwen3_moe, gpt_oss — add .to(device=q.device) on cos/sin before rotary apply ## MoE meta guards (weight-free needs shape-only meta placeholders) - mixtral_moe: QEffMixtralExperts.__qeff_init__ with meta guard - grok_1: QEffGrok1MoeBlock.__qeff_init__ with meta guard - qwen3_moe: QEffQwen3MoeExperts — add explicit meta guard - glm4_moe: QEffGlm4MoeMoE.__qeff_init__ — meta guard + rename parameters from self.all_gate_proj → self.experts.gate_proj so ONNX initializer names match MoEFusedExpertSplitCheckpointTransform output keys Signed-off-by: Amar <amarshar@qti.qualcomm.com>
Infrastructure:
- modeling_qeff.py: reorder example_inputs/dynamic_shapes by forward()
signature so torch.export binds dynamic_shapes correctly; fixes
position_ids aliased as past_key.0 for GPT-style models causing
duplicate graph inputs and compiler errors
- modeling_qeff.py: restore weight_spec_path on ONNX cache hit, guard
use_weight_free_export without dynamo, skip SplitTensorsTransform,
embed weight_spec.json as ONNX metadata, symlink prepared checkpoint
- core.py: add dtype suffix to prepared checkpoint dir, prune fake
initializers after meta-device ONNX export
- modeling_auto.py: auto-enable dynamo=True with use_weight_free_export
Model wrapper fixes:
- codegen: replace embed_positions side-effect with .to(dtype, device)
to fix float32 subfunction type mismatch for weight-free export
- glm4_moe: cast logits with .float() to fix 2x runtime buffer mismatch;
fix rotary cos/sin device sync, _rotary_dim attribute, router device,
MoE expert placeholder dtype
- granitemoe/mixtral/phi3/qwen3_moe: fix rotary cos/sin device/dtype
for weight-free meta-device tracing
- mixtral: refactor expert forward to use derived params directly so
fused gate_up_proj meta tensor never appears as ONNX initializer;
fix MoE expert placeholder dtype
- continuous_batching.py: rename use_dynamo=True to dynamo=True to
prevent invalid -use-dynamo compiler flag
Tests:
- Add tests/weight_free/ with 75 tests across 21 model families:
ONNX structure, HF PT == ORT parity, CB export, transform unit tests
- New assertions: assert_unique_graph_input_names and
assert_no_int64_kv_cache_inputs guard position_ids aliasing regression
Signed-off-by: Amar <amarshar@qti.qualcomm.com>
…nit test related fixes Signed-off-by: Amar <amarshar@qti.qualcomm.com>
Code quality (modeling_qeff.py): - Remove redundant DYNAMO_CUSTOM_OP_TABLE import; register custom ops in customop/__init__.py so torch.ops.qefficient.* are always available - Factor duplicated input-reordering and dynamo kwargs setup into reorder_inputs_by_signature() and build_dynamo_export_kwargs() in export_utils.py (shared by dynamo and weight-free paths) - Fix weight_spec_path type: both assignments now str() consistently - Expose public embed_weight_spec_as_metadata() in core.py; base class no longer imports private _upsert_metadata_prop across modules - Wrap symlink creation in try/except OSError with logger.warning - Promote use_weight_free_export to named param in _compile() signature - Add Tuple[Optional[Dict], Any] return type to _export_via_weightfree Architecture (export backends): - Create QEfficient/exporter/onnx_exporter.py with export_via_legacy, export_via_dynamo, export_via_weightfree as standalone functions; _export() calls them directly without wrapper methods - Pass pre-computed dynamic_shapes to export_via_weightfree (computed once by export_wrapper) — removes redundant convert_dynamic_axes_to_ dynamic_shapes call inside the function; signatures are now parallel - Keep _export_layerwise in base class (VLM-specific branching needs a separate refactor with proper override hooks) Examples: - Add examples/dynamo/causal_lm/weight_free_inference.py — clean dedicated weight-free example following basic_dynamo_inference.py - Update basic_inference.py: add --use-weight-free-export flag with meta-device model path when flag is set - Update continuous_batching.py: add --use-weight-free-export flag - Move weight_free/_runner.py → scripts/weight_free_runner.py - Remove weight_free/export_compile_infer.py (covered by basic_inference) - Remove weight_free/continuous_batching.py (merged into existing CB) - Remove weight_free/__init__.py Deprecation: - Add DeprecationWarning for layerwise export in from_pretrained(), export(), and compile() — points users to use_weight_free_export=True Signed-off-by: Amar <amarshar@qti.qualcomm.com>
Move .to(device=hidden_states.device) to the sin/cos slice at the point of indexing in each model's forward method, rather than inside the rotary embedding helper function. Previously the helper checked cos.device != q.device and moved the full sin_cached/cos_cached buffer. With this change only the small indexed slice [position_ids].unsqueeze(1) is moved, which is cheaper and keeps the original buffer on its init device. Also removes the redundant device check from qeff_apply_rotary_pos_emb in each model — the caller is now responsible for device alignment. Models updated: llama, mistral, qwen2, qwen3, qwen3_moe, falcon, gemma, gemma2, granite, granitemoe, olmo2, phi3, gpt_oss, glm4_moe, mixtral Signed-off-by: Amar <amarshar@qti.qualcomm.com>
Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>
Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>
runtime, and export orchestration modules. Rename checkpoint transform helpers to checkpoint_transforms.py to align with existing transform naming patterns. Make checkpoint preparation conditional, avoid mutating source checkpoint directories during .bin to safetensors conversion, and remove script-level profiler imports from exporter code. Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>
Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>
…runner
- weight_free_runner.py: cb_ccl mode now enables CCL context lengths
(previously only "ccl" did); fix use_dynamo -> dynamo kwarg name in
export/compile calls; widen default CCL list to {1024, 2048, ctx_len};
add --skip_compile to isolate export/checkpoint-prep timing.
- basic_dynamo_inference.py: default num_devices to 2 when no device
group is given.
- tests/weight_free: promote gpt_oss from xfail to a real test target
using openai/gpt-oss-20b, matching exported layer count to the HF
model for ORT parity; add on-QAIC hardware parity tests (weight-free
generate smoke test, HF-PT vs QAIC parity, weight-free vs legacy
QAIC parity).
Signed-off-by: amarshar <amarshar@qti.qualcomm.com>
Not needed by tests/weight_free/test_export.py or test_on_qaic.py, which build and drive QEff models directly. Keeping the runner's benchmarking changes local rather than in this PR. Signed-off-by: amarshar <amarshar@qti.qualcomm.com>
Only used as a local benchmarking helper by scripts/compare_weightfree.py (not tracked in this repo). Not referenced by any test or production code. Signed-off-by: amarshar <amarshar@qti.qualcomm.com>
…g dtype Thread enable_proxy through basic_dynamo_inference.py's load_qeff_model so proxy mode can be exercised via --enable-proxy alongside --use-weight-free-export. QeffProxyEmbedding.forward hardcoded its output to float32 via .float(), which breaks non-float32 models (e.g. float16 Llama) with a dtype mismatch in downstream attention matmuls. Cast to the original embedding weight's dtype instead, since the in-place class swap preserves the original nn.Embedding parameters. Signed-off-by: amarshar <amarshar@qti.qualcomm.com>
test_weight_free_export_ort_parity, test_weight_free_hw_hf_parity, and test_weight_free_vs_legacy_qaic_parity currently fail for gpt_oss. Mark them xfail, matching the existing gpt_oss xfail pattern in test_continuous_batching.py, so the suite reflects the known gap instead of reporting hard failures. Signed-off-by: amarshar <amarshar@qti.qualcomm.com>
Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com>
ochougul
merged commit Aug 7, 2026
366a956
into
quic:quic-enable-weightfree-for-causallm
17 checks passed
quic-amitraj
added a commit
that referenced
this pull request
Aug 7, 2026
Signed-off-by: Amar <amarshar@qti.qualcomm.com> Signed-off-by: Amit Raj <amitraj@qti.qualcomm.com> Co-authored-by: Amit Raj <amitraj@qti.qualcomm.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
3 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.
No description provided.