vllm 0.4.1 compatibility

[mgerstgrasser]

vllm 0.4.1 refactors model loading, this PR updates OpenRLHF to be compatible with the changes.

Monkey patching the model weights iterator is no longer necessary, and some imports change slightly.

[hijkzzz]

Thanks for you MR @wuxibin89 will review it.

[wuxibin89]

@mgerstgrasser Great that we no longer need this nasty monkey patch for hf_model_weights_iterator.

According to https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py#L392, the update_weight function should also be adjusted

    def update_weight(self, name, dtype, shape, empty_cache=False):
        ...
        if vllm.__version__ < "0.4.1":
            self.model_runner.model.load_weights(model_name_or_path={name: weight})
        else:
            self.model_runner.model.load_weights(weights=[(name,weight)])

[hijkzzz]

We could also update the dockerfile:

https://github.com/OpenLLMAI/OpenRLHF/blob/7f8239249a3b982f757b867cddc6808b158263b5/dockerfile/Dockerfile#L18

RUN pip install vllm==0.4.1

[hijkzzz]

I revert to this MR because it can't pass the test

[hijkzzz]

Fixed in https://github.com/OpenLLMAI/OpenRLHF/pull/283

[mgerstgrasser]

@hijkzzz @wuxibin89

the update_weight function should also be adjusted

Ah, apologies for the oversight!

Is this already fixed now? Or is there anything left I can help with?