[Bug] Cannot build `--quantization q3f16_1` model

[isaac621]

🐛 Bug

Meet an error while building q3f16_1 vicuna-7b-v1.5, the same error also occur when I try to building q3f16_1 for other models

To Reproduce

python3 -m mlc_llm.build --hf-path lmsys/vicuna-7b-v1.5 --target iphone --max-seq-len 4096 --quantization q3f16_1

Error Message

Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/isaaclui/Desktop/oursky-project/mlc-llm/mlc_llm/build.py", line 13, in <module>
    main()
  File "/Users/isaaclui/Desktop/oursky-project/mlc-llm/mlc_llm/build.py", line 10, in main
    core.build_model_from_args(parsed_args)
  File "/Users/isaaclui/Desktop/oursky-project/mlc-llm/mlc_llm/core.py", line 653, in build_model_from_args
    build(mod, args)
  File "/Users/isaaclui/Desktop/oursky-project/mlc-llm/mlc_llm/core.py", line 514, in build
    mod_deploy = dl.ApplyDefaultSchedule(  # pylint: disable=not-callable
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/ir/transform.py", line 238, in __call__
    return _ffi_transform_api.RunPass(self, mod)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/_ffi/_cython/./packed_func.pxi", line 332, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 263, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 252, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 182, in tvm._ffi._cy3.core.CHECK_CALL
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/_ffi/base.py", line 476, in raise_last_ffi_error
    raise py_err
  File "tvm/_ffi/_cython/./packed_func.pxi", line 56, in tvm._ffi._cy3.core.tvm_callback
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/ir/transform.py", line 307, in _pass_func
    return inst.transform_module(mod, ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/dlight/base/transform.py", line 64, in transform_module
    sch = _apply_rules(func, target, self.rules, tunable=False)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/dlight/base/transform.py", line 80, in _apply_rules
    space = rule.apply(func, target, tunable)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/dlight/gpu/gemv.py", line 191, in apply
    is_inner_reduction = normalize(sch, block_info)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/isaaclui/.pyenv/versions/3.11.5/lib/python3.11/site-packages/tvm/dlight/gpu/gemv.py", line 122, in normalize
    is_inner_reduction = iter_to_info[inner_axis].kind == "R"
                         ~~~~~~~~~~~~^^^^^^^^^^^^
KeyError: IterSplit(IterMark(v1, extent=T.int64(4096)), lower_factor=T.int64(1), extent=T.int64(4096), scale=T.int64(1))

[junrushao]

I'm not exactly sure what happened but it seems your codebase has been a bit outdated. Would you like to use the latest TVM Unity? https://llm.mlc.ai/docs/install/tvm.html#option-1-prebuilt-package

[isaac621]

@junrushao I have already reinstall the TVM Unity via the command python3 -m pip install --pre --force-reinstall -f https://mlc.ai/wheels mlc-ai-nightly, but same error is happening when I try to compile the model to q3f16_1.

FYI, it works well when I compile the model to q4f16_0

[tqchen]

cc @vinx13 can you help looking into this one?

[isaac621]

@vinx13 Would like to know any update on this? Thanks a lot 🙇‍♂️

[vinx13]

the regression is caused by apache/tvm#15665 I’m still working on a fix, as a workaround you can also revert it for now

[isaac621]

@vinx13 For the workaround, should I build the tvm from the source code with commit before this apache/tvm#15665?

The same error occur after I install the tvm from the revert build and run the model compile script.

[vinx13]

I applied the fix on the head of unity branch

[isaac621]

I applied the fix on the head of unity branch

Cannot see the fix commit 🤔

[junrushao]

@vinx13 could you also cherry pick this fix to the mlc branch?

[vinx13]

I sent a fix in apache/tvm#15881, it will be cherry-picked once merged

[isaac621]

@vinx13 Have installed the tvm based on your fix commit on dc53a6c29. But the same problem still occur when I try to compile the q3f16. Anything have I done wrong?

[vinx13]

I ran python3 -m mlc_llm.build --hf-path lmsys/vicuna-7b-v1.5 --target iphone --max-seq-len 4096 --quantization q3f16_1 on unity branch as of today and can compile successfully. Maybe you can check your installation and make sure you are using the tvm you installed. The fix merged to unity now, you can also install from https://llm.mlc.ai/docs/install/tvm.html#option-1-prebuilt-package

[isaac621]

It works on my side now! Thx you a lot @vinx13 🔥