[Bug] TVMError: Check failed: (partial_update) is false: Key "vocab_size" not found.

[KiwiHana]

NV 3060 laptop
cuda12.1

$ python build.py --model llama2/Llama-2-7b-chat-hf --target cuda --quantization q4f16_1 --artifact-path "./dist" --use-cache 0

$ python benchmark.py

from mlc_chat import ChatModule

cm = ChatModule(model="./dist/Llama-2-7b-chat-hf-q4f16_1/params/", lib_path="./dist/Llama-2-7b-chat-hf-q4f16_1/Llama-2-7b-chat-hf-q4f16_1-cuda.so")

output = cm.benchmark_generate("What's the meaning of life?", generate_length=256)
print(f"Generated text:\n{output}\n")
print(f"Statistics: {cm.stats()}")

$ python benchmark.py
Traceback (most recent call last):
  File "/home/kiwi/llm/mlc-llm/benchmark.py", line 9, in <module>
    cm = ChatModule(model="./dist/Llama-2-7b-chat-hf-q4f16_1/params/", lib_path="./dist/Llama-2-7b-chat-hf-q4f16_1/Llama-2-7b-chat-hf-q4f16_1-cuda.so",device="vulkan")
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kiwi/miniconda3/envs/mlc-py11/lib/python3.11/site-packages/mlc_chat/chat_module.py", line 603, in __init__
    self._reload(self.lib_path, self.model_path, user_chat_config_json_str)
  File "/home/kiwi/miniconda3/envs/mlc-py11/lib/python3.11/site-packages/mlc_chat/chat_module.py", line 784, in _reload
    self._reload_func(lib, model_path, app_config_json)
  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 "/home/kiwi/miniconda3/envs/mlc-py11/lib/python3.11/site-packages/tvm/_ffi/base.py", line 476, in raise_last_ffi_error
    raise py_err
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 1183, in mlc::llm::LLMChatModule::GetFunction(tvm::runtime::String const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 468, in mlc::llm::LLMChat::Reload(tvm::runtime::String, tvm::runtime::String, tvm::runtime::String)
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 448, in mlc::llm::LLMChat::LoadJSONOverride(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, bool)
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 373, in mlc::llm::LLMChat::LoadJSONOverride(picojson::value const&, bool)
tvm._ffi.base.TVMError: Traceback (most recent call last):
  3: mlc::llm::LLMChatModule::GetFunction(tvm::runtime::String const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const
        at /workspace/mlc-llm/cpp/llm_chat.cc:1183
  2: mlc::llm::LLMChat::Reload(tvm::runtime::String, tvm::runtime::String, tvm::runtime::String)
        at /workspace/mlc-llm/cpp/llm_chat.cc:468
  1: mlc::llm::LLMChat::LoadJSONOverride(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, bool)
        at /workspace/mlc-llm/cpp/llm_chat.cc:448
  0: mlc::llm::LLMChat::LoadJSONOverride(picojson::value const&, bool)
        at /workspace/mlc-llm/cpp/llm_chat.cc:373
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 373
TVMError: Check failed: (partial_update) is false: Key "vocab_size" not found.

[MasterJH5574]

HI @KiwiHana, could you check "./dist/Llama-2-7b-chat-hf-q4f16_1/params/mlc-chat-config.json" to see if vocab_size is a field there? Now we do generate the config with vocab_size:

mlc-llm/mlc_llm/core.py

Line 484 in bd28bd2

config["vocab_size"] = vocab_size

If vocab_size does not exist in your mlc-chat-config.json, could you update mlc-llm to the latest commit and try again?

[KiwiHana]

Yes, it work. But NV 3060 6GB has Error out of memory, when VRAM use 4.7/6.0 GB

$ python benchmark.py
Traceback (most recent call last):
  File "/mnt/c/kiwi/llm/llama2/mlc-llm/benchmark.py", line 9, in <module>
    cm = ChatModule(model="./dist/Llama-2-7b-chat-hf-q4f16_1/params/", lib_path="./dist/Llama-2-7b-chat-hf-q4f16_1/Llama-2-7b-chat-hf-q4f16_1-cuda.so")
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/kiwi/miniconda3/envs/mlc-py11/lib/python3.11/site-packages/mlc_chat/chat_module.py", line 603, in __init__
    self._reload(self.lib_path, self.model_path, user_chat_config_json_str)
  File "/home/kiwi/miniconda3/envs/mlc-py11/lib/python3.11/site-packages/mlc_chat/chat_module.py", line 784, in _reload
    self._reload_func(lib, model_path, app_config_json)
  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 "/home/kiwi/miniconda3/envs/mlc-py11/lib/python3.11/site-packages/tvm/_ffi/base.py", line 476, in raise_last_ffi_error
    raise py_err
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 1183, in mlc::llm::LLMChatModule::GetFunction(tvm::runtime::String const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const
  File "/workspace/mlc-llm/cpp/llm_chat.cc", line 493, in mlc::llm::LLMChat::Reload(tvm::runtime::String, tvm::runtime::String, tvm::runtime::String)
tvm.error.InternalError: Traceback (most recent call last):
  12: mlc::llm::LLMChatModule::GetFunction(tvm::runtime::String const&, tvm::runtime::ObjectPtr<tvm::runtime::Object> const&)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*) const
        at /workspace/mlc-llm/cpp/llm_chat.cc:1183
  11: mlc::llm::LLMChat::Reload(tvm::runtime::String, tvm::runtime::String, tvm::runtime::String)
        at /workspace/mlc-llm/cpp/llm_chat.cc:493
  10: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeClosurePacked(tvm::runtime::ObjectRef const&, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  9: tvm::runtime::PackedFuncObj::Extractor<tvm::runtime::PackedFuncSubObj<tvm::runtime::relax_vm::VirtualMachineImpl::GetClosureInternal(tvm::runtime::String const&, bool)::{lambda(tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)#1}> >::Call(tvm::runtime::PackedFuncObj const*, tvm::runtime::TVMArgs, tvm::runtime::TVMRetValue*)
  8: tvm::runtime::relax_vm::VirtualMachineImpl::InvokeBytecode(long, std::vector<tvm::runtime::TVMRetValue, std::allocator<tvm::runtime::TVMRetValue> > const&)
  7: tvm::runtime::relax_vm::VirtualMachineImpl::RunLoop()
  6: tvm::runtime::relax_vm::VirtualMachineImpl::RunInstrCall(tvm::runtime::relax_vm::VMFrame*, tvm::runtime::relax_vm::Instruction)
  5: _ZN3tvm7runtime13PackedFun
  4: tvm::runtime::TypedPackedFunc<tvm::runtime::relax_vm::AttentionKVCache (tvm::runtime::NDArray, tvm::runtime::ShapeTuple, int)>::AssignTypedLambda<tvm::runtime::relax_vm::AttentionKVCache (*)(tvm::runtime::NDArray, tvm::runtime::ShapeTuple, int)>(tvm::runtime::relax_vm::AttentionKVCache (*)(tvm::runtime::NDArray, tvm::runtime::ShapeTuple, int), std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >)::{lambda(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*)#1}::operator()(tvm::runtime::TVMArgs const&, tvm::runtime::TVMRetValue*) const
  3: tvm::runtime::relax_vm::AttentionKVCache::Create(tvm::runtime::NDArray, tvm::runtime::ShapeTuple, int)
  2: tvm::runtime::NDArray::Empty(tvm::runtime::ShapeTuple, DLDataType, DLDevice, tvm::runtime::Optional<tvm::runtime::String>)
  1: tvm::runtime::CUDADeviceAPI::AllocDataSpace(DLDevice, unsigned long, unsigned long, DLDataType)
  0: _ZN3tvm7runtime6deta
  File "/workspace/tvm/src/runtime/cuda/cuda_device_api.cc", line 129
InternalError: Check failed: (e == cudaSuccess || e == cudaErrorCudartUnloading) is false: CUDA: out of memory

[MasterJH5574]

Glad that it works. The default max sequence length for Llama2 is 4096 (as set in the Llama2 config https://huggingface.co/meta-llama/Llama-2-7b-chat-hf/blob/main/config.json#L12). To save memory, you can try compile with the argument --max-seq-len 2048 (or even smaller) in the command. For example,

python build.py --model llama2/Llama-2-7b-chat-hf --target cuda --quantization q4f16_1 --artifact-path "./dist" --use-cache 0 --max-seq-len 1024

[KiwiHana]

Yes, you are right! one more question, what is the performance of mlc-llm in the following log, is it decode?
Statistics: prefill: tok/s, decode: tok/s

[MasterJH5574]

prefill: xxx tok/s means the performance of the prefill stage (taking the user input prompt), and decode: xxx tok/s means the performance of the decode stage (keep generating new tokens).

[KiwiHana]

prefill: xxx tok/s means the performance of the prefill stage (taking the user input prompt), and decode: xxx tok/s means the performance of the decode stage (keep generating new tokens).

Thank you so much!