Repository navigation
Commit b1e9c8e
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-5daa19ed42d41 parent d734346 commit b1e9c8e
15 files changed
Lines changed: 1243 additions & 91 deletions
File tree
- src/mobius
- components
- integrations
- transformers
- models
- weights
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
47 | 47 | | |
48 | 48 | | |
49 | 49 | | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| 59 | + | |
| 60 | + | |
| 61 | + | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
| 79 | + | |
| 80 | + | |
| 81 | + | |
| 82 | + | |
| 83 | + | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
| 90 | + | |
| 91 | + | |
| 92 | + | |
| 93 | + | |
| 94 | + | |
| 95 | + | |
| 96 | + | |
| 97 | + | |
| 98 | + | |
| 99 | + | |
| 100 | + | |
| 101 | + | |
| 102 | + | |
| 103 | + | |
50 | 104 | | |
51 | 105 | | |
52 | 106 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
43 | 43 | | |
44 | 44 | | |
45 | 45 | | |
| 46 | + | |
46 | 47 | | |
47 | 48 | | |
48 | 49 | | |
| |||
1415 | 1416 | | |
1416 | 1417 | | |
1417 | 1418 | | |
| 1419 | + | |
1418 | 1420 | | |
1419 | 1421 | | |
1420 | 1422 | | |
| |||
1432 | 1434 | | |
1433 | 1435 | | |
1434 | 1436 | | |
| 1437 | + | |
| 1438 | + | |
| 1439 | + | |
1435 | 1440 | | |
1436 | 1441 | | |
1437 | 1442 | | |
1438 | 1443 | | |
1439 | 1444 | | |
1440 | 1445 | | |
1441 | | - | |
| 1446 | + | |
1442 | 1447 | | |
1443 | 1448 | | |
1444 | 1449 | | |
1445 | | - | |
| 1450 | + | |
1446 | 1451 | | |
1447 | 1452 | | |
1448 | 1453 | | |
| 1454 | + | |
| 1455 | + | |
| 1456 | + | |
1449 | 1457 | | |
1450 | | - | |
1451 | | - | |
| 1458 | + | |
1452 | 1459 | | |
1453 | 1460 | | |
1454 | 1461 | | |
1455 | 1462 | | |
1456 | 1463 | | |
1457 | 1464 | | |
1458 | 1465 | | |
1459 | | - | |
| 1466 | + | |
| 1467 | + | |
| 1468 | + | |
| 1469 | + | |
| 1470 | + | |
1460 | 1471 | | |
1461 | 1472 | | |
1462 | 1473 | | |
1463 | 1474 | | |
1464 | 1475 | | |
1465 | 1476 | | |
1466 | 1477 | | |
1467 | | - | |
| 1478 | + | |
| 1479 | + | |
| 1480 | + | |
1468 | 1481 | | |
1469 | 1482 | | |
1470 | | - | |
| 1483 | + | |
1471 | 1484 | | |
1472 | 1485 | | |
1473 | 1486 | | |
1474 | 1487 | | |
1475 | | - | |
| 1488 | + | |
1476 | 1489 | | |
| 1490 | + | |
| 1491 | + | |
| 1492 | + | |
| 1493 | + | |
1477 | 1494 | | |
1478 | 1495 | | |
1479 | 1496 | | |
| |||
1629 | 1646 | | |
1630 | 1647 | | |
1631 | 1648 | | |
1632 | | - | |
| 1649 | + | |
| 1650 | + | |
| 1651 | + | |
| 1652 | + | |
| 1653 | + | |
1633 | 1654 | | |
1634 | 1655 | | |
1635 | 1656 | | |
| |||
1643 | 1664 | | |
1644 | 1665 | | |
1645 | 1666 | | |
1646 | | - | |
| 1667 | + | |
| 1668 | + | |
| 1669 | + | |
| 1670 | + | |
| 1671 | + | |
| 1672 | + | |
1647 | 1673 | | |
1648 | 1674 | | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1440 | 1440 | | |
1441 | 1441 | | |
1442 | 1442 | | |
| 1443 | + | |
| 1444 | + | |
| 1445 | + | |
| 1446 | + | |
| 1447 | + | |
| 1448 | + | |
| 1449 | + | |
| 1450 | + | |
| 1451 | + | |
| 1452 | + | |
| 1453 | + | |
| 1454 | + | |
| 1455 | + | |
| 1456 | + | |
| 1457 | + | |
| 1458 | + | |
| 1459 | + | |
| 1460 | + | |
| 1461 | + | |
| 1462 | + | |
| 1463 | + | |
| 1464 | + | |
| 1465 | + | |
| 1466 | + | |
| 1467 | + | |
| 1468 | + | |
| 1469 | + | |
| 1470 | + | |
| 1471 | + | |
| 1472 | + | |
| 1473 | + | |
| 1474 | + | |
| 1475 | + | |
| 1476 | + | |
| 1477 | + | |
| 1478 | + | |
| 1479 | + | |
| 1480 | + | |
| 1481 | + | |
| 1482 | + | |
| 1483 | + | |
| 1484 | + | |
| 1485 | + | |
| 1486 | + | |
| 1487 | + | |
| 1488 | + | |
| 1489 | + | |
| 1490 | + | |
| 1491 | + | |
| 1492 | + | |
| 1493 | + | |
| 1494 | + | |
| 1495 | + | |
| 1496 | + | |
| 1497 | + | |
| 1498 | + | |
| 1499 | + | |
| 1500 | + | |
| 1501 | + | |
| 1502 | + | |
| 1503 | + | |
| 1504 | + | |
| 1505 | + | |
| 1506 | + | |
| 1507 | + | |
| 1508 | + | |
| 1509 | + | |
| 1510 | + | |
| 1511 | + | |
| 1512 | + | |
| 1513 | + | |
| 1514 | + | |
| 1515 | + | |
| 1516 | + | |
| 1517 | + | |
| 1518 | + | |
| 1519 | + | |
| 1520 | + | |
| 1521 | + | |
| 1522 | + | |
| 1523 | + | |
| 1524 | + | |
| 1525 | + | |
| 1526 | + | |
| 1527 | + | |
| 1528 | + | |
| 1529 | + | |
| 1530 | + | |
| 1531 | + | |
| 1532 | + | |
| 1533 | + | |
| 1534 | + | |
| 1535 | + | |
| 1536 | + | |
| 1537 | + | |
| 1538 | + | |
| 1539 | + | |
| 1540 | + | |
| 1541 | + | |
| 1542 | + | |
| 1543 | + | |
1443 | 1544 | | |
1444 | 1545 | | |
1445 | 1546 | | |
| |||
0 commit comments