[Bug]: Llama3 VocabParallelEmbedding error when loading

[gelim]

Your current environment

PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 10.5.0-1ubuntu1~22.04) 10.5.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.4
Libc version: glibc-2.35

Python version: 3.11.0rc1 (main, Aug 12 2022, 10:02:14) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-92-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration:
GPU 0: Tesla V100-SXM2-32GB
GPU 1: Tesla V100-SXM2-32GB
GPU 2: Tesla V100-SXM2-32GB
GPU 3: Tesla V100-SXM2-32GB
GPU 4: Tesla V100-SXM2-32GB
GPU 5: Tesla V100-SXM2-32GB
GPU 6: Tesla V100-SXM2-32GB
GPU 7: Tesla V100-SXM2-32GB

Nvidia driver version: 535.154.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             12
On-line CPU(s) list:                0-11
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Bronze 3104 CPU @ 1.70GHz
CPU family:                         6
Model:                              85
Thread(s) per core:                 1
Core(s) per socket:                 6
Socket(s):                          2
Stepping:                           4
CPU max MHz:                        1700.0000
CPU min MHz:                        800.0000
BogoMIPS:                           3400.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single pti intel_ppin ssbd mba ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm arat pln pts hwp hwp_act_window hwp_pkg_req pku ospke md_clear flush_l1d arch_capabilities
Virtualization:                     VT-x
L1d cache:                          384 KiB (12 instances)
L1i cache:                          384 KiB (12 instances)
L2 cache:                           12 MiB (12 instances)
L3 cache:                           16.5 MiB (2 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-2,6-8
NUMA node1 CPU(s):                  3-5,9-11
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Mitigation; PTE Inversion; VMX conditional cache flushes, SMT disabled
Vulnerability Mds:                  Mitigation; Clear CPU buffers; SMT disabled
Vulnerability Meltdown:             Mitigation; PTI
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT disabled
Vulnerability Retbleed:             Mitigation; IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; IBRS, IBPB conditional, STIBP disabled, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; Clear CPU buffers; SMT disabled

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] nvidia-nccl-cu12==2.20.5
[pip3] pyzmq==26.2.0
[pip3] torch==2.4.0
[pip3] torchvision==0.19.0
[pip3] transformers==4.45.2
[pip3] triton==3.0.0
[conda] Could not collect
ROCM Version: Could not collect
Neuron SDK Version: N/A
Aphrodite Version: 0.6.2.post1
Aphrodite Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0    GPU1    GPU2    GPU3    GPU4    GPU5    GPU6    GPU7    CPU Affinity    NUMA Affinity   GPU NUMA ID
GPU0     X      NV1     NV1     NV2     NV2     SYS     SYS     SYS     0-2,6-8 0               N/A
GPU1    NV1      X      NV2     NV1     SYS     NV2     SYS     SYS     0-2,6-8 0               N/A
GPU2    NV1     NV2      X      NV2     SYS     SYS     NV1     SYS     0-2,6-8 0               N/A
GPU3    NV2     NV1     NV2      X      SYS     SYS     SYS     NV1     0-2,6-8 0               N/A
GPU4    NV2     SYS     SYS     SYS      X      NV1     NV1     NV2     3-5,9-11        1               N/A
GPU5    SYS     NV2     SYS     SYS     NV1      X      NV2     NV1     3-5,9-11        1               N/A
GPU6    SYS     SYS     NV1     SYS     NV1     NV2      X      NV2     3-5,9-11        1               N/A
GPU7    SYS     SYS     SYS     NV1     NV2     NV1     NV2      X      3-5,9-11        1               N/A

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks```

🐛 Describe the bug

Helo, when trygin to run GGUF of Llama-3.1-8B, LLama-3.2-3B, I always get an error at start time:

$ CUDA_VISIBLE_DEVICES=6 python -m aphrodite.endpoints.openai.api_server --model /models/Llama-3.2-3B-Instruct-Q8_0.gguf -tp 1 --dev
ice cuda --api-keys xxxx--served-model-name llama3.2-3b --host 127.0.0.1 --port 8080
INFO:     Multiprocessing frontend to use ipc:///tmp/3c0ca7f0-d165-4b79-b9a9-f6f7d9f5060e for RPC Path.
INFO:     Started engine process with PID 414584
WARNING:  gguf quantization is not fully optimized yet. The speed can be slower than non-quantized models.
WARNING:  The model has a long context length (131072). This may cause OOM errors during the initial memory profiling phase, or result in low performance due to small KV cache space. Consider setting
--max-model-len to a smaller value.
INFO:     -------------------------------------------------------------------------------------
INFO:     Initializing Aphrodite Engine (v0.6.2.post1 commit 4d3d819) with the following config:
INFO:     Model = '/models/Llama-3.2-3B-Instruct-Q8_0.gguf'
INFO:     DataType = torch.float16
INFO:     Model Load Format = \<LoadFormat.GGUF: 'gguf'>
INFO:     Tensor Parallel Size = 1
INFO:     Pipeline Parallel Size = 1
INFO:     Disable Custom All-Reduce = False
INFO:     Quantization Format = 'gguf'
INFO:     Context Length = 131072
INFO:     Enforce Eager Mode = False
INFO:     Prefix Caching = False
INFO:     Device = device(type='cuda')
INFO:     Guided Decoding Backend = DecodingConfig(guided_decoding_backend='outlines')
INFO:     -------------------------------------------------------------------------------------
The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization.
The tokenizer class you load from this checkpoint is 'BartTokenizer'.
The class this function is called from is 'LlamaTokenizerFast'.
You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama_fast.LlamaTokenizerFast'>. This is expected, and simply means that the `legacy` (previous) behavior will b
e used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explai
ned in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message.
INFO:     Cannot use FlashAttention-2 backend for Volta and Turing GPUs.
INFO:     Using XFormers backend.
/home/mgeli-local/aphrodite/lib/python3.11/site-packages/xformers/ops/fmha/flash.py:211: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we w
ill remove `torch.library.impl_abstract` in a future version of PyTorch.                                                                                                                                           @torch.library.impl_abstract("xformers_flash::flash_fwd")
/home/mgeli-local/aphrodite/lib/python3.11/site-packages/xformers/ops/fmha/flash.py:344: FutureWarning: `torch.library.impl_abstract` was renamed to `torch.library.register_fake`. Please use that instead; we w
ill remove `torch.library.impl_abstract` in a future version of PyTorch.
  @torch.library.impl_abstract("xformers_flash::flash_bwd")
INFO:     Loading model /models/Llama-3.2-3B-Instruct-Q8_0.gguf...
INFO:     Cannot use FlashAttention-2 backend for Volta and Turing GPUs.
INFO:     Using XFormers backend.
Process SpawnProcess-1:
Traceback (most recent call last):
  File "/usr/lib/python3.11/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/usr/lib/python3.11/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/endpoints/openai/rpc/server.py", line 209, in run_rpc_server
    server = AsyncEngineRPCServer(async_engine_args, rpc_path)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/endpoints/openai/rpc/server.py", line 24, in __init__
    self.engine = AsyncAphrodite.from_engine_args(async_engine_args)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/engine/async_aphrodite.py", line 601, in from_engine_args
    engine = cls(
             ^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/engine/async_aphrodite.py", line 510, in __init__
    self.engine = self._init_engine(*args, **kwargs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/engine/async_aphrodite.py", line 694, in _init_engine
    return engine_class(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/engine/aphrodite_engine.py", line 261, in __init__
    self.model_executor = executor_class(
                          ^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/executor/executor_base.py", line 45, in __init__
    self._init_executor()
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/executor/gpu_executor.py", line 36, in _init_executor
    self.driver_worker.load_model()
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/task_handler/worker.py", line 153, in load_model
    self.model_runner.load_model()
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/task_handler/model_runner.py", line 888, in load_model
    self.model = get_model(model_config=self.model_config,
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/modeling/model_loader/__init__.py", line 20, in get_model
    return loader.load_model(model_config=model_config,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/modeling/model_loader/loader.py", line 1033, in load_model
    model = _initialize_model(model_config, self.load_config,
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/modeling/model_loader/loader.py", line 167, in _initialize_model
    return build_model(
           ^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/modeling/model_loader/loader.py", line 152, in build_model
    return model_class(config=hf_config,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/aphrodite/modeling/models/llama.py", line 427, in __init__
    self.lm_head.weight = self.model.embed_tokens.weight
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgeli-local/aphrodite/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1729, in __getattr__
    raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
AttributeError: 'VocabParallelEmbedding' object has no attribute 'weight'

Running something like Yi-1.5-9B-Chat-Q4_K_M.gguf loads fine.
I am not sure where to look at, as it seems the VocabParallelEmbedding is a new feature of Llama3 (see https://medium.com/@vlad_96627/llama3-vs-llama2-whats-new-65f47fc9b4ec) and you claim you are supporting it.