Skip to content

Commit b1e9c8e

Browse files
titaiwangmsCopilot
andauthored
Add projection-specific INT2/INT3/INT4/INT8 QMoE export (#764)
## 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 - 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: ```bash 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. ```bash lintrunner src/mobius/_model_package.py src/mobius/_model_package_test.py \ src/mobius/_weight_utils.py src/mobius/integrations/_weight_loading.py \ src/mobius/integrations/transformers/_builder.py src/mobius/models/granite.py \ src/mobius/models/int3_qmoe_test.py src/mobius/models/llada.py \ src/mobius/models/moe.py src/mobius/models/qwen35.py \ src/mobius/models/qwen35_test.py src/mobius/weights/_adapters.py git diff --check ``` **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. - INT2 runtime/CI qualification is separate in #765 and is not included in this PR. --------- Signed-off-by: titaiwang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: fba8a408-cd63-49f0-b8b4-5daa19ed42d4
1 parent d734346 commit b1e9c8e

15 files changed

Lines changed: 1243 additions & 91 deletions

File tree

‎README.md‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,60 @@ across **40+ task types** and **100+ reusable components**.
4747

4848
See the [model documentation](https://onnxruntime.github.io/mobius/models/index.html) for the complete list.
4949

50+
## Projection-specific QMoE export
51+
52+
Olive's integer-affine, fused, K-last expert checkpoints can independently
53+
use INT2, INT3, INT4, or INT8 for FC1 and FC2, with a supported model-wide
54+
INT4 fallback for non-expert projections. All 16 bit-width pairs are accepted
55+
when their layouts satisfy the checks below; no experimental flag is required. INT3
56+
uses the proposed raw QMoE contract from `microsoft/onnxruntime#32657`.
57+
58+
```python
59+
from mobius import build
60+
61+
package = build("./local-olive-qwen3-moe", dtype="f16")
62+
package.save("./qmoe-export")
63+
```
64+
65+
Use an Olive model-wide INT4 configuration with exact source-name overrides for
66+
`model.layers.<i>.mlp.experts.gate_up_proj` and
67+
`model.layers.<i>.mlp.experts.down_proj`. The generic MoE adapter, including Qwen3-MoE,
68+
validates fused checkpoint dimensions and binds packed bytes without
69+
requantization. INT3 is symmetric unsigned offset-4, row-local LSB-first
70+
`uint8`; FC1 gate/up rows and scales are interleaved together, while FC2
71+
retains its K-last down-projection layout. Explicit all-4 INT3 zero points are
72+
validated and omitted in the graph.
73+
74+
Both projections require a shared power-of-two block size of at least 16,
75+
dividing hidden and intermediate sizes. Pairs containing INT3 require block
76+
size 32/64/128, symmetric INT3, and FP16/BF16 activations. For these pairs both
77+
FC scale dtypes must match the activations. Every projection's scales must be
78+
finite and nonnegative. When the adapter supplies the activation dtype,
79+
scales must remain finite after conversion. Non-INT3 layouts retain the
80+
standard weight loader's dtype conversion.
81+
The legacy uniform INT4 path receives the same numeric scale checks.
82+
INT2, INT4, and INT8 projections retain their existing zero-point
83+
semantics; asymmetric zero points must be packed `uint8` sidecars.
84+
Dense INT3 `MatMulNBits`, embeddings, other source layouts, and schema-unapproved
85+
tails are not enabled. Keep non-expert projections at supported precisions.
86+
87+
`build()` defers INT3-layout FC1 interleaving to initializer serialization,
88+
using uncached lazy tensors instead of retaining a second complete FC1 payload.
89+
The original checkpoint is still loaded eagerly; this bounds additional
90+
interleaving storage, not total checkpoint memory. Direct `preprocess_weights()`
91+
calls retain eager behavior unless an output `weight_transforms` mapping is
92+
provided. Pass that mapping to `ModelPackage.apply_weights()` with the returned
93+
weights to preserve the deferred transformations.
94+
95+
**Export support is not a runtime qualification matrix.** INT3 QMoE nodes
96+
carry a draft-format/runtime-warning marker; they do not invent an
97+
approved schema version. Execution of every pair depends on the target ORT
98+
version and execution provider. Tests cover bytes, interleaving, binding, and external
99+
data save/reload, not INT3 kernel execution, model quality, or speed. Final
100+
schema/version review, a complete quantization provenance manifest, and
101+
checkpoint-to-ORT parity remain necessary before production use. This is not
102+
completion of the runtime plan's P5 acceptance gate.
103+
50104
## Installation
51105

52106
```bash

‎src/mobius/_model_package.py‎

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343

4444
from mobius._export_report import ComponentExportReport
4545
from mobius._optimizations import fold_initializers_after_weights
46+
from mobius._weight_utils import WeightTransform
4647
from mobius.adapters import (
4748
AdapterArtifact,
4849
AdapterServiceOptions,
@@ -1415,6 +1416,7 @@ def apply_weights(
14151416
prefix_map: dict[str, str] | None = None,
14161417
*,
14171418
fold_constants: bool = True,
1419+
weight_transforms: dict[str, WeightTransform] | None = None,
14181420
) -> None:
14191421
"""Apply weights from a state dict across component models.
14201422
@@ -1432,48 +1434,63 @@ def apply_weights(
14321434
Weights whose name starts with a prefix are applied to the
14331435
named component (with the prefix stripped). Unmatched weights
14341436
are applied to all components.
1437+
weight_transforms: Shape-preserving transforms keyed by names in
1438+
``state_dict`` (the preprocessed names, before prefix routing),
1439+
evaluated without caching during serialization.
14351440
"""
14361441
applied: set[str] = set()
14371442

14381443
if len(self.data) == 1:
14391444
# Single component — apply all weights directly
14401445
model = next(iter(self.data.values()))
1441-
applied = _apply_weights_to_model(model, state_dict)
1446+
applied = _apply_weights_to_model(model, state_dict, weight_transforms)
14421447
elif prefix_map is None:
14431448
# No routing info — try every weight against every model
14441449
for model in self.data.values():
1445-
applied |= _apply_weights_to_model(model, state_dict)
1450+
applied |= _apply_weights_to_model(model, state_dict, weight_transforms)
14461451
else:
14471452
# Route by prefix
14481453
routed: dict[str, dict[str, torch.Tensor]] = {name: {} for name in self.data}
1454+
routed_transforms: dict[str, dict[str, WeightTransform]] = {
1455+
name: {} for name in self.data
1456+
}
14491457
unmatched: dict[str, torch.Tensor] = {}
1450-
# Track original HF names for weights that get stripped
1451-
stripped_to_original: dict[str, str] = {}
1458+
stripped_to_original: dict[tuple[str, str], str] = {}
14521459

14531460
for weight_name, tensor in state_dict.items():
14541461
matched = False
14551462
for prefix, component in prefix_map.items():
14561463
if weight_name.startswith(prefix):
14571464
stripped = weight_name[len(prefix) :].lstrip(".")
14581465
routed[component][stripped] = tensor
1459-
stripped_to_original[stripped] = weight_name
1466+
if weight_transforms is not None and weight_name in weight_transforms:
1467+
routed_transforms[component][stripped] = weight_transforms[
1468+
weight_name
1469+
]
1470+
stripped_to_original[(component, stripped)] = weight_name
14601471
matched = True
14611472
break
14621473
if not matched:
14631474
unmatched[weight_name] = tensor
14641475

14651476
for component_name, component_weights in routed.items():
14661477
applied_stripped = _apply_weights_to_model(
1467-
self.data[component_name], component_weights
1478+
self.data[component_name],
1479+
component_weights,
1480+
routed_transforms[component_name],
14681481
)
14691482
for s in applied_stripped:
1470-
applied.add(stripped_to_original.get(s, s))
1483+
applied.add(stripped_to_original.get((component_name, s), s))
14711484

14721485
# Try unmatched weights against all models
14731486
if unmatched:
14741487
for model in self.data.values():
1475-
applied |= _apply_weights_to_model(model, unmatched)
1488+
applied |= _apply_weights_to_model(model, unmatched, weight_transforms)
14761489

1490+
if weight_transforms is not None:
1491+
unused = set(weight_transforms) - applied
1492+
if unused:
1493+
raise ValueError(f"Weight transforms were not bound: {sorted(unused)!r}.")
14771494
_log_weight_mapping(state_dict, applied)
14781495

14791496
if not fold_constants:
@@ -1629,7 +1646,11 @@ def _log_weight_mapping(
16291646
)
16301647

16311648

1632-
def _apply_weights_to_model(model: ir.Model, state_dict: dict[str, torch.Tensor]) -> set[str]:
1649+
def _apply_weights_to_model(
1650+
model: ir.Model,
1651+
state_dict: dict[str, torch.Tensor],
1652+
weight_transforms: dict[str, WeightTransform] | None = None,
1653+
) -> set[str]:
16331654
"""Apply weights to a single model (internal helper).
16341655
16351656
Uses :func:`_assign_weight` from ``_weight_loading`` for shape
@@ -1643,6 +1664,11 @@ def _apply_weights_to_model(model: ir.Model, state_dict: dict[str, torch.Tensor]
16431664
if name not in model.graph.initializers:
16441665
continue
16451666

1646-
_assign_weight(model.graph.initializers[name], tensor, name)
1667+
_assign_weight(
1668+
model.graph.initializers[name],
1669+
tensor,
1670+
name,
1671+
transform=weight_transforms.get(name) if weight_transforms is not None else None,
1672+
)
16471673
applied.add(name)
16481674
return applied

‎src/mobius/_model_package_test.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1440,6 +1440,107 @@ def test_save_preserves_public_policy_symbols(self, tmp_path):
14401440

14411441

14421442
class TestModelPackageApplyWeights:
1443+
@pytest.mark.parametrize("second_has_initializer", [False, True])
1444+
def test_weight_transform_accounting_separates_component_local_names(
1445+
self, second_has_initializer
1446+
):
1447+
values = {}
1448+
models = {}
1449+
for component in ("text", "vision"):
1450+
graph = ir.Graph([], [], nodes=[], name=component)
1451+
if component == "text" or second_has_initializer:
1452+
value = ir.Value(
1453+
name="weight",
1454+
type=ir.TensorType(ir.DataType.FLOAT),
1455+
shape=ir.Shape([2, 3]),
1456+
const_value=ir.tensor(torch.zeros(2, 3)),
1457+
)
1458+
graph.register_initializer(value)
1459+
values[component] = value
1460+
models[component] = ir.Model(graph, ir_version=10)
1461+
package = ModelPackage(models)
1462+
weights = {"text.weight": torch.zeros(2, 3), "vision.weight": torch.ones(2, 3)}
1463+
transforms = {
1464+
"text.weight": lambda tensor: tensor + 2,
1465+
"vision.weight": lambda tensor: tensor + 3,
1466+
}
1467+
if not second_has_initializer:
1468+
with pytest.raises(ValueError, match=r"not bound.*vision\.weight"):
1469+
package.apply_weights(
1470+
weights,
1471+
prefix_map={"text.": "text", "vision.": "vision"},
1472+
weight_transforms={"vision.weight": transforms["vision.weight"]},
1473+
fold_constants=False,
1474+
)
1475+
else:
1476+
package.apply_weights(
1477+
weights,
1478+
prefix_map={"text.": "text", "vision.": "vision"},
1479+
weight_transforms=transforms,
1480+
fold_constants=False,
1481+
)
1482+
for component in ("text", "vision"):
1483+
actual = torch.from_numpy(values[component].const_value.numpy())
1484+
expected = transforms[component + ".weight"](weights[component + ".weight"])
1485+
assert torch.equal(actual, expected)
1486+
1487+
@pytest.mark.parametrize("routed", [False, True])
1488+
def test_weight_transforms_are_lazy_and_follow_prefix_routing(self, routed):
1489+
value = ir.Value(
1490+
name="weight",
1491+
type=ir.TensorType(ir.DataType.FLOAT),
1492+
shape=ir.Shape([2, 3]),
1493+
const_value=ir.tensor(torch.zeros(2, 3)),
1494+
)
1495+
graph = ir.Graph([], [], nodes=[], name="text")
1496+
graph.register_initializer(value)
1497+
package = ModelPackage(
1498+
{"text": ir.Model(graph, ir_version=10), "other": _make_simple_model()}
1499+
)
1500+
calls = []
1501+
source = torch.arange(6).reshape(2, 3).float()
1502+
1503+
def transform(tensor):
1504+
calls.append(tensor.data_ptr())
1505+
return tensor + 1
1506+
1507+
name = "text.weight" if routed else "weight"
1508+
package.apply_weights(
1509+
{name: source},
1510+
prefix_map={"text.": "text"} if routed else None,
1511+
weight_transforms={name: transform},
1512+
fold_constants=False,
1513+
)
1514+
assert calls == []
1515+
assert isinstance(value.const_value, ir.LazyTensor)
1516+
assert not value.const_value.cache
1517+
for count in (1, 2):
1518+
assert torch.equal(torch.from_numpy(value.const_value.numpy()), source + 1)
1519+
assert len(calls) == count
1520+
1521+
def test_unbound_weight_transform_raises(self):
1522+
package = ModelPackage({"model": _make_simple_model()})
1523+
with pytest.raises(ValueError, match="Weight transforms were not bound"):
1524+
package.apply_weights({}, weight_transforms={"missing": lambda tensor: tensor})
1525+
1526+
def test_weight_transform_changed_shape_raises(self):
1527+
value = ir.Value(
1528+
name="weight",
1529+
type=ir.TensorType(ir.DataType.FLOAT),
1530+
shape=ir.Shape([2, 3]),
1531+
const_value=ir.tensor(torch.zeros(2, 3)),
1532+
)
1533+
graph = ir.Graph([], [], nodes=[], name="model")
1534+
graph.register_initializer(value)
1535+
package = ModelPackage({"model": ir.Model(graph, ir_version=10)})
1536+
package.apply_weights(
1537+
{"weight": torch.zeros(2, 3)},
1538+
weight_transforms={"weight": lambda tensor: tensor.reshape(3, 2)},
1539+
fold_constants=False,
1540+
)
1541+
with pytest.raises(ValueError, match="changed shape"):
1542+
value.const_value.numpy()
1543+
14431544
def test_single_component(self):
14441545
config = make_config()
14451546
module = CausalLMModel(config)

0 commit comments

Comments
 (0)