RewardModel.from_pretrained() loads redundant weights (incurs extra ~30GB of RAM)

[angie-chen55]

Hi,

Whenever a saved RewardModel is loaded via RewardModel.from_pretrained(model_path, flash_attn=True, fp16=False, bf16=True, low_cpu_mem_usage=True), it downloads the entire sharded checkpoint (https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/modeling_utils.py#L2876), which is already ~30GB because it contains all the weights of the reward model, including both the backbone model and the reward head. It then calls RewardModel.__init__() (via this line https://github.com/huggingface/transformers/blob/v4.33.2/src/transformers/modeling_utils.py#L2966), which loads all the weights of the backbone model (SFT10K, another ~30GB). Surely loading a pretrained model shouldn't require loading the backbone model weights twice?

Thanks!