Enable weightfree for causallm1 - #1220
Open
amarquic wants to merge 2 commits into
Open
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>
amarquic
force-pushed
the
enable_weightfree_for_causallm1
branch
from
July 27, 2026 18:48
c309de1 to
c6828f5
Compare
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.
Summary
Adds weight-free ONNX export for causal LM models on top of the dynamo export infrastructure.
Weight-free export builds the model on meta device (no weights in RAM), exports ONNX graph structure only, and lets the QAIC compiler load
weights directly from the original safetensors checkpoint at compile time via
weight_spec.json.New files
QEfficient/exporter/weight_free/— core weight-free export, weight spec format, checkpoint prep pipeline (MoE expert stacking, dtypeconversion, .bin→safetensors auto-convert)
examples/text_generation/weight_free/— single-model smoke test, CB weight-free smoke testKey changes
modeling_qeff.py—use_weight_free_exportflag; reorderexample_inputs/dynamic_shapesbyforward()signature sotorch.exportbinds them correctly; skip
SplitTensorsTransform; embedweight_spec.jsonas ONNX metadata; symlink prepared checkpoint next to ONNXmodeling_auto.py—use_weight_free_exportwired throughexport()andcompile(); auto-enabledynamo=Trueonnx_transforms.py— disable unsupported transforms for dynamo pathexport_utils.py— separate dynamo/TorchScript subfunction setup pathstorch_patches.py—temporarily_enable/disable_nested_compile_regions()core.py— dtype suffix in prepared checkpoint dir; prune fake initializers after meta-device exportModel fixes
codegen— fixembed_positionsdevice/dtype for weight-free subfunction tracingglm4_moe—logits.float()to fix runtime buffer mismatch; rotary cos/sin device sync;_rotary_dimattribute; MoE expert placeholderdtype
granitemoe,mixtral,phi3,qwen3_moe— rotary cos/sin device/dtype sync for meta-device tracing; MoE expert placeholder dtypecontinuous_batching.py— fixuse_dynamo=True→dynamo=TrueTests
pip install -e .[test] pytest tests/weight_free/ -m "not on_qaic" -n auto -v