You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Add projection-specific QMoE export for Olive fused expert checkpoints.
FC1 and FC2 independently support INT2, INT3, INT4, and INT8: all 16 pairs
are accepted subject to format validation, without an experimental opt-in.
This is export support, not runtime qualification. INT3 follows the proposed
raw-weight direction in microsoft/onnxruntime#32657 and consumes the native Torch
INT3 checkpoint format introduced by microsoft/Olive#2716. INT3 output retains
draft-format/runtime-warning metadata; this PR does not provide INT3 ORT kernels
or claim that released ORT can execute it.
Changes
Resolve FC1/FC2 widths independently instead of using an explicit pair whitelist.
This intentionally includes (2,3), (3,2), (3,8), (8,3), (2,2), and (2,8), in addition to the other valid pairs.
Remove the experimental INT3 flag and opt-in rejection. Export requires an
explicitly quantized checkpoint with compatible projection configurations;
ordinary unquantized models do not become INT3 automatically.
Require exact Olive expert overrides and a supported model-wide INT4 fallback
for projection-specific layouts. Preserve the existing uniform INT4 path.
Preserve row-local, LSB-first packed unsigned INT3 codes with symmetric zero
point 4. For pairs containing INT3, interleave FC1 gate/up payloads lazily during
serialization in public HF builds; direct preprocessing remains eager by
default. Otherwise retain graph-side interleaving. Preserve FC2's K-last layout
without requantization.
Validate shared power-of-two group sizes, aligned dimensions, sidecar shapes,
packed zero points, and finite nonnegative scales. Reject scale conversion
overflow when the adapter supplies the activation dtype. INT3-containing pairs
additionally require block size 32/64/128 and FP16/BF16 activation/scale
agreement. Non-INT3 layouts retain the standard loader's dtype conversion.
Apply the same numeric scale checks to the legacy uniform INT4 path without
imposing the projection-specific shape ABI. Rejecting previously accepted
negative, NaN, infinite, or activation-cast-overflowing scales is an intentional
validation tightening.
Bind deferred FC1 weights, scales, and asymmetric zero points through uncached
lazy tensors, rather than retaining transformed copies for every layer.
Keep transform accounting component-scoped when prefix stripping produces
identical local initializer names.
Preserve INT2/INT4/INT8 zero-point semantics and validate optional explicit
offset-4 INT3 zero points before omitting them from the graph.
Emit projection-width attributes where needed and retain INT3 draft-format
metadata without assigning an unapproved schema version.
Cover all 16 pairs across FP16/BF16, independent packing/dequantization oracles,
CPU execution of graph-side interleaving, invalid scales, local-checkpoint
public builds, and external-data save/reload.
Pass the model activation dtype through all three Qwen3.5-MoE preprocessing
paths: text, VL global quantization, and VL decoder component quantization.
Add 57 regression cases covering all seven INT3-containing pairs across
FP16/BF16, incompatible dtypes, and non-INT3 scale conversion behavior.
Validation
Latest offline CPU validation includes the merged main, Qwen3.5-MoE dtype
propagation, uniform scale checks, deferred serialization, and component routing:
PYTHONPATH="$PWD/src" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
TOKENIZERS_PARALLELISM=false PYTHONDONTWRITEBYTECODE=1 CUDA_VISIBLE_DEVICES='' \
timeout 300 /home/titaiwang/miniconda3/envs/dev/bin/python -m pytest \
--rootdir="$PWD" -c "$PWD/pyproject.toml" \
src/mobius/models/int3_qmoe_test.py src/mobius/models/qwen35_test.py \
src/mobius/models/moe_test.py \
src/mobius/_weight_utils_test.py src/mobius/_flags_test.py \
src/mobius/components/_moe_test.py \
src/mobius/_model_package_test.py::TestModelPackageApplyWeights \
src/mobius/weights/_adapters_test.py \
src/mobius/integrations/_weight_loading_test.py \
src/mobius/integrations/transformers/_builder_test.py \
-q -m 'not integration and not arch_validation and not golden and not generation' \
--tb=short
Passed: 728 tests, 1 deselected, in 52.36 seconds. Coverage includes
distinct per-layer payloads, source-storage aliasing, release of transformed
tensors between callbacks, repeated saves, asymmetric FC1 zero points, and
positive/negative component-local-name collision regressions.
Both passed. Ruff formatting and checking also passed for the same 12 Python
files. Scoped mypy reported 77 errors outside changed hunks; no project-wide
type-check success or fully established clean baseline is claimed.
Draft boundaries
All 16 pairs are exporter-supported layouts, not a runtime-qualified matrix.
Execution depends on the target ORT version and execution provider.
Source checkpoint loading remains eager. Uncached lazy interleaving bounds
additional FC1 payload retention, not total checkpoint or serialization peak
memory. No full-checkpoint RSS or concurrent-writer qualification was run.
Direct callers opting into deferred preprocessing must forward the populated
transform map when binding the returned weights. Explicit adapters and older
preprocessing hooks retain their existing behavior.
No INT3 ORT kernel execution, checkpoint-to-runtime logits parity, real-model
quality, or performance qualification was run. The CPU interleave tests
exercise standard ONNX graph operations, not QMoE kernels.
Dense INT3 MatMulNBits, quantized embeddings, unsupported source layouts, and
schema-unapproved tails are not enabled.
Final QMoE schema/version agreement, complete quantization provenance, and
runtime qualification remain follow-up gates. This does not complete the
end-to-end plan's P5 acceptance gate.
Allow all FC1/FC2 precision combinations in {2, 3, 4, 8} without an experimental opt-in while retaining layout validation and the INT3 runtime qualification warning. Add scale conversion guards and independent graph interleave and quantization convention coverage.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: fba8a408-cd63-49f0-b8b4-5daa19ed42d4
Signed-off-by: titaiwang <titaiwang@microsoft.com>
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Eager FC1 interleaving can nearly double host-memory usage for production MoE checkpoints, and the documented scale-validation contract has an uncovered exception.
This states that every projection's scales receive finite/nonnegative and conversion checks, but the new checks only run when projection_layouts is nonempty. The legacy uniform (4,4) pair has requires_projection_preprocessing == False, continues through preprocess_olive_weights, and therefore accepts NaN/negative/overflowing scales. Either apply the validation to that path too or document the uniform-INT4 exception so the advertised 16-pair contract is accurate.
Avoid retaining duplicate FC1 weights during preprocessing
src/mobius/_weight_utils.py:1279
This interleave materializes a contiguous copy of every FC1 packed weight (and scale tensor), while all original tensors remain referenced by state_dict and all transformed copies accumulate in renamed_experts. For production MoE checkpoints, FC1 is several gigabytes, so preprocessing can nearly double the routed-expert payload and make otherwise viable exports run out of host memory. Please transform these tensors through the streaming weight pipeline (or otherwise release each source tensor as soon as its transformed replacement is owned) instead of eagerly retaining both complete layouts.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
The draft INT3 operator contract and cross-cutting lazy serialization changes warrant final human review.
0 open findings
🧠 Review effort: Balanced
This branch has not been deployed
No deployments
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
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
Add projection-specific QMoE export for Olive fused expert checkpoints.
FC1 and FC2 independently support INT2, INT3, INT4, and INT8: all 16 pairs
are accepted subject to format validation, without an experimental opt-in.
This is export support, not runtime qualification. INT3 follows the proposed
raw-weight direction in microsoft/onnxruntime#32657 and consumes the native Torch
INT3 checkpoint format introduced by microsoft/Olive#2716. INT3 output retains
draft-format/runtime-warning metadata; this PR does not provide INT3 ORT kernels
or claim that released ORT can execute it.
Changes
This intentionally includes
(2,3),(3,2),(3,8),(8,3),(2,2), and(2,8), in addition to the other valid pairs.explicitly quantized checkpoint with compatible projection configurations;
ordinary unquantized models do not become INT3 automatically.
for projection-specific layouts. Preserve the existing uniform INT4 path.
point 4. For pairs containing INT3, interleave FC1 gate/up payloads lazily during
serialization in public HF builds; direct preprocessing remains eager by
default. Otherwise retain graph-side interleaving. Preserve FC2's K-last layout
without requantization.
packed zero points, and finite nonnegative scales. Reject scale conversion
overflow when the adapter supplies the activation dtype. INT3-containing pairs
additionally require block size 32/64/128 and FP16/BF16 activation/scale
agreement. Non-INT3 layouts retain the standard loader's dtype conversion.
imposing the projection-specific shape ABI. Rejecting previously accepted
negative, NaN, infinite, or activation-cast-overflowing scales is an intentional
validation tightening.
lazy tensors, rather than retaining transformed copies for every layer.
Keep transform accounting component-scoped when prefix stripping produces
identical local initializer names.
offset-4 INT3 zero points before omitting them from the graph.
metadata without assigning an unapproved schema version.
CPU execution of graph-side interleaving, invalid scales, local-checkpoint
public builds, and external-data save/reload.
paths: text, VL global quantization, and VL decoder component quantization.
Add 57 regression cases covering all seven INT3-containing pairs across
FP16/BF16, incompatible dtypes, and non-INT3 scale conversion behavior.
Validation
Latest offline CPU validation includes the merged
main, Qwen3.5-MoE dtypepropagation, uniform scale checks, deferred serialization, and component routing:
Passed: 728 tests, 1 deselected, in 52.36 seconds. Coverage includes
distinct per-layer payloads, source-storage aliasing, release of transformed
tensors between callbacks, repeated saves, asymmetric FC1 zero points, and
positive/negative component-local-name collision regressions.
Both passed. Ruff formatting and checking also passed for the same 12 Python
files. Scoped mypy reported 77 errors outside changed hunks; no project-wide
type-check success or fully established clean baseline is claimed.
Draft boundaries
Execution depends on the target ORT version and execution provider.
additional FC1 payload retention, not total checkpoint or serialization peak
memory. No full-checkpoint RSS or concurrent-writer qualification was run.
transform map when binding the returned weights. Explicit adapters and older
preprocessing hooks retain their existing behavior.
quality, or performance qualification was run. The CPU interleave tests
exercise standard ONNX graph operations, not QMoE kernels.
MatMulNBits, quantized embeddings, unsupported source layouts, andschema-unapproved tails are not enabled.
runtime qualification remain follow-up gates. This does not complete the
end-to-end plan's P5 acceptance gate.
included in this PR.