Issue with running script

[Viroslav]

Hello, after all following your code and installing vllm_uncertainty I'm trying to run a code. In forward_step:

(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/AINTJOB/Research/RAG/SeaKR/vllm_uncertainty/vllm/model_executor/models/llama.py", line 292, in forward
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     hidden_states, residual = layer(
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/miniconda3/envs/seakr_new/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     return self._call_impl(*args, **kwargs)
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/miniconda3/envs/seakr_new/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     return forward_call(*args, **kwargs)
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/AINTJOB/Research/RAG/SeaKR/vllm_uncertainty/vllm/model_executor/models/llama.py", line 225, in forward
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     hidden_states = self.input_layernorm(hidden_states)
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/miniconda3/envs/seakr_new/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     return self._call_impl(*args, **kwargs)
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/miniconda3/envs/seakr_new/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     return forward_call(*args, **kwargs)
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/AINTJOB/Research/RAG/SeaKR/vllm_uncertainty/vllm/model_executor/layers/layernorm.py", line 60, in forward
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     ops.rms_norm(
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]   File "/media/ssd-3t/Vir_Ven/AINTJOB/Research/RAG/SeaKR/vllm_uncertainty/vllm/_custom_ops.py", line 120, in rms_norm
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148]     vllm_ops.rms_norm(out, input, weight, epsilon)
(RayWorkerWrapper pid=3634342) ERROR 11-07 23:41:39 worker_base.py:148] NameError: name 'vllm_ops' is not defined

I got an error of 'vllm_ops' undefined. Then I go to _custom_ops.py: And I found the fix with import

Instead of importing from vllm._C I did from vllm.model_executor.layers import ops as vllm_ops But that doesn't have a method rms_norm in line:

# layer norm ops
def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
             epsilon: float) -> None:
    vllm_ops.rms_norm(out, input, weight, epsilon)
    

I took a look into original vllm and they use torch instead:

# layer norm ops
def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
             epsilon: float) -> None:
    torch.ops._C.rms_norm(out, input, weight, epsilon)
    ```
    
    But if i change it to torch, it goes nuts and a lot of errors come in. What should I do?