feat: mx integration by chichun-charlie-liu · Pull Request #110 · foundation-model-stack/fms-model-optimizer

chichun-charlie-liu · 2025-04-29T19:08:13Z

Description of the change

Support mx formats by integrating mx package's functionalities. Triggered by mx_ prefix in config's qw_mode, qa_mode,...

Related issue number

part #105 but need two more PRs

use diff/patch file to resolve microxcaling's dependency/installation issue
add supports to fp8e4m3, e5m2 format for scales.
(see Issue 105 for more details)

How to verify the PR

New unit tests are added, including config and toy model tests.

Was the PR tested

I have added >=1 unit test(s) for every new method I have added.
I have ensured all unit tests pass

Signed-off-by: cliu-us <cliu@us.ibm.com>

…. previously has to be all or none mx Signed-off-by: cliu-us <cliu@us.ibm.com>

…-existing quantizer "mx_abc", force nbits=32 for QBmm init to avoid getting quantizers. QBmmMX simply relies on qa_mode, not nbits Signed-off-by: cliu-us <cliu@us.ibm.com>

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

feat: Mx impl brandon

Signed-off-by: chichun-charlie-liu <57839396+chichun-charlie-liu@users.noreply.github.com>

tharapalanivel

Did a quick review and looks like lint and some of the tests are failing sadly but great PR @chichun-charlie-liu and the docs particularly look great 😄 Thank you!

Signed-off-by: cliu-us <cliu@us.ibm.com>

…ings Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

fix: Setting mx_specs outside qconfig_init

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

…bugs Signed-off-by: cliu-us <cliu@us.ibm.com>

Signed-off-by: chichun-charlie-liu <57839396+chichun-charlie-liu@users.noreply.github.com>

Signed-off-by: cliu-us <cliu@us.ibm.com>

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Signed-off-by: cliu-us <cliu@us.ibm.com>

tharapalanivel · 2025-05-27T21:54:40Z

@@ -46,3 +46,4 @@ data*_train/
 data*_test/
 act_scales/
 examples/


I'm not sure why we're ignoring examples dir since we actually have example docs there, can we get rid of it?

adjusted, now we will ignore generated files such as .pt, .json, .safetensors, .arrow but allow others, such as .md and .py

tharapalanivel · 2025-05-27T21:56:25Z

    if qcfg.get("QBmm"):  # see Note 4
        # Local
-        from fms_mo.modules import QBmm
+        QBmm = qcfg["mapping"]["matmul_or_bmm"]


Should this be lower case variable name?

This is a class pointer disguised as a variable and used as QBmm(args) below. I think it is fine as is.

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Signed-off-by: cliu-us <cliu@us.ibm.com>

chichun-charlie-liu and others added 20 commits March 27, 2025 18:45

toy exam under examples/MX

53f9857

include mx Linear class directly for convenience

a59010e

Signed-off-by: cliu-us <cliu@us.ibm.com>

minor clean up

b6b1d4f

Signed-off-by: cliu-us <cliu@us.ibm.com>

[pyproject] Added MX optional-dependency

027eb2f

[qcfg] Changed mapping init and added mx path in qconfig_init

02a9c2d

[QLinear] Added QLinearMX stub

2bd90e9

[qcfg] Changed qcfg[mapping] structure

cc681ee

[mx] Updates to ffn_tmp.py

62962a3

clean-up, now simply use qa_mode and qw_mode to trigger mx

a79ebbf

Signed-off-by: cliu-us <cliu@us.ibm.com>

first mx test with "dq" using qa_/qw_mode=mx_fp8_e4m3

8bf30ac

Signed-off-by: cliu-us <cliu@us.ibm.com>

add QbmmMX support

ac03469

Signed-off-by: cliu-us <cliu@us.ibm.com>

draft example readme for mx

7f6ea36

Signed-off-by: cliu-us <cliu@us.ibm.com>

adjust mapping def, allow qbmm to use non-mx while linear is using mx…

613e523

…. previously has to be all or none mx Signed-off-by: cliu-us <cliu@us.ibm.com>

bug fix, QBmmMX init was calling QBmmMX init which tries to get a non…

9ea77da

…-existing quantizer "mx_abc", force nbits=32 for QBmm init to avoid getting quantizers. QBmmMX simply relies on qa_mode, not nbits Signed-off-by: cliu-us <cliu@us.ibm.com>

feat: Added mx_specs to qconfig_init and check_config

91c20cd

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

feat: Added skippable mx_specs unit tests

8b1d026

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

fix: Added QLinearMX and QBmmMX to quantized_modules

a161b9a

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

chore: format qconfig_utils

9ffa0b1

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

test: minor fixes to test_mx and supporting functions

a9a57b3

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Merge pull request #1 from BrandonGroth/mx_impl_brandon

764d7bc

feat: Mx impl brandon

chichun-charlie-liu requested review from BrandonGroth, andrea-fasoli, kcirred, nwang-ibm and tharapalanivel as code owners April 29, 2025 19:08

github-actions Bot added the feat label Apr 29, 2025

Merge branch 'main' into mx_impl

45dd501

Signed-off-by: chichun-charlie-liu <57839396+chichun-charlie-liu@users.noreply.github.com>

tharapalanivel requested changes Apr 29, 2025

View reviewed changes

lint and ruff fix

9bab835

Signed-off-by: cliu-us <cliu@us.ibm.com>

chichun-charlie-liu marked this pull request as draft April 29, 2025 20:34

chichun-charlie-liu and others added 9 commits May 14, 2025 17:31

rename mx example py, markdown update is still WIP

d7dfc27

Signed-off-by: cliu-us <cliu@us.ibm.com>

markdown update is still WIP

3446413

Signed-off-by: cliu-us <cliu@us.ibm.com>

fix bug in config check as QBmmMX is optional

dd17104

Signed-off-by: cliu-us <cliu@us.ibm.com>

feat: Added qmodel_prep mx_specs hook for default config with MX sett…

6e4431c

…ings Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

test: Added MX test for default config + qmodel_prep

83d2b06

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

fix: Added MX vocabulary for spellcheck util

d05401f

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

chore: Fixed MX vocab errors

7daae3b

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

fix: Changed check_config check for MX mapping

5694a5a

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

Merge pull request #3 from BrandonGroth/mx_impl_brandon

308e5e2

fix: Setting mx_specs outside qconfig_init

BrandonGroth reviewed May 23, 2025

View reviewed changes

Comment thread examples/MX/README.md Outdated

doc: Updated MX examples/readme for patch install

e5eb0ed

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

BrandonGroth approved these changes May 23, 2025

View reviewed changes

update mx readme.md, fix simple_mx_exam.py together with a few minor …

ab060ca

…bugs Signed-off-by: cliu-us <cliu@us.ibm.com>

chichun-charlie-liu marked this pull request as ready for review May 27, 2025 17:11

Merge branch 'main' into mx_impl

9b9b777

Signed-off-by: chichun-charlie-liu <57839396+chichun-charlie-liu@users.noreply.github.com>

BrandonGroth reviewed May 27, 2025

View reviewed changes

Comment thread examples/MX/README.md

chichun-charlie-liu and others added 3 commits May 27, 2025 17:51

add a microscaling ref in readme

2cb9aca

Signed-off-by: cliu-us <cliu@us.ibm.com>

fix: Patched qconfig_init and get_unwanted_defaults QBmm

7e7a0fa

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

fix: Updated spellcheck list for examples/mx readme

f220b17

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

BrandonGroth force-pushed the mx_impl branch from cc22b42 to f220b17 Compare May 27, 2025 19:22

BrandonGroth approved these changes May 27, 2025

View reviewed changes

fix: Fixed additional Qbmm mapping in test_saveconfig

026a3ea

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

BrandonGroth force-pushed the mx_impl branch from a153102 to 026a3ea Compare May 27, 2025 20:05

minor updates on README

b00287f

Signed-off-by: cliu-us <cliu@us.ibm.com>

tharapalanivel requested changes May 27, 2025

View reviewed changes

BrandonGroth and others added 2 commits May 28, 2025 01:02

chore: Cleanup of extra comments

b5f0eaa

Signed-off-by: Brandon Groth <brandon.m.groth@gmail.com>

adjust gitignore to allow changes in py and md files

4e0de3f

Signed-off-by: cliu-us <cliu@us.ibm.com>

tharapalanivel approved these changes May 28, 2025

View reviewed changes

chichun-charlie-liu merged commit 05bb442 into foundation-model-stack:main May 28, 2025
12 checks passed

chichun-charlie-liu deleted the mx_impl branch May 28, 2025 16:34

Conversation

chichun-charlie-liu commented Apr 29, 2025

Description of the change

Related issue number

How to verify the PR

Was the PR tested

Uh oh!

tharapalanivel left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tharapalanivel May 27, 2025

Choose a reason for hiding this comment

Uh oh!

chichun-charlie-liu May 28, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tharapalanivel May 27, 2025

Choose a reason for hiding this comment

Uh oh!

BrandonGroth May 28, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tharapalanivel left a comment •

edited

Loading