Using LLaMA in reward model training

[YingHH1]

Hi,

I have encounter the TypeError: LlamaModel.forward() got an unexpected keyword argument 'head_mask' error when training the LLaMA-7B model in step 2 reward model training.

I was wondering if the head_mask is used at all in training the reward model?

Also, is there a quick fix for this error?

Many thanks

[jose77]

The LLaMA model does not have the "head_mask" argument, which you can remove in the utils/model/reward_model.py.

[YingHH1]

The LLaMA model does not have the "head_mask" argument, which you can remove in the utils/model/reward_model.py.

Thanks! So the reward model training doesn't use the head_mask at all then?

[qinqinqaq]

The LLaMA model does not have the "head_mask" argument, which you can remove in the utils/model/reward_model.py.

When i remove “head_mask”

I get error:
Traceback (most recent call last):
File "main.py", line 357, in
main()
File "main.py", line 311, in main
reward_score, acc = evaluation_reward(rm_model, eval_dataloader)
File "main.py", line 257, in evaluation_reward
outputs = model(**batch)
File "/home/kemove/anaconda3/envs/deepspeed/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "/home/kemove/anaconda3/envs/deepspeed/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/home/kemove/anaconda3/envs/deepspeed/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1695, in forward
loss = self.module(*inputs, **kwargs)
File "/home/kemove/anaconda3/envs/deepspeed/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1208, in _call_impl
result = forward_call(*input, **kwargs)
File "/home/kemove/fkb/project/deepspeed/DeepSpeedExamples/applications/DeepSpeed-Chat/training/utils/model/reward_model.py", line 202, in forward
rewards = self.v_head(hidden_states).squeeze(-1)
File "/home/kemove/anaconda3/envs/deepspeed/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1208, in _call_impl
result = forward_call(*input, **kwargs)
File "/home/kemove/anaconda3/envs/deepspeed/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (4096x32000 and 4096x1)

[YingHH1]

When I remove the head_mask, I am met with the same divergence_id assertion error as in #338

[qinqinqaq]

I modified the following code and it worked
# self.config.n_embd = self.config.hidden_size if hasattr(
# self.config, "hidden_size") else self.config.n_embd
# todo
self.config.n_embd = self.config.vocab_size if hasattr(
self.config, "vocab_size") else self.config.n_embd

[zhouliang-yu]

I modified the following code and it worked # self.config.n_embd = self.config.hidden_size if hasattr( # self.config, "hidden_size") else self.config.n_embd # todo self.config.n_embd = self.config.vocab_size if hasattr( self.config, "vocab_size") else self.config.n_embd

Thx for your solution, after modifying the code according to your suggestion, I finally make it also work on qwen