error when launching reward model training with the weights produced in SFT stage

[javismiles]

I followed your instructions for the SFT training and I got the .pt trained weights,
then I run:

python train_rm.py -b 2 -n experiment_name -p "./runs/sft_javSFT2_202405021344/sft_javSFT2_202405021344_step4000.pt

and I get this error in line 87 of train_rm.py

File "/xxxxxxxxxxxxxxxxxxxx/python3.9/site-packages/torch/nn/modules/module.py", line 2152, in load_state_dict
raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for GPT:
Missing key(s) in state_dict: "transformer.decoder_blocks.0.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.0.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.0.mmsa.output_projection.lora_A", "transformer.decoder_blocks.0.mmsa.output_projection.lora_B", "transformer.decoder_blocks.0.ffn.fc1.lora_A", "transformer.decoder_blocks.0.ffn.fc1.lora_B", "transformer.decoder_blocks.0.ffn.fc2.lora_A", "transformer.decoder_blocks.0.ffn.fc2.lora_B", "transformer.decoder_blocks.1.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.1.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.1.mmsa.output_projection.lora_A", "transformer.decoder_blocks.1.mmsa.output_projection.lora_B", "transformer.decoder_blocks.1.ffn.fc1.lora_A", "transformer.decoder_blocks.1.ffn.fc1.lora_B", "transformer.decoder_blocks.1.ffn.fc2.lora_A", "transformer.decoder_blocks.1.ffn.fc2.lora_B", "transformer.decoder_blocks.2.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.2.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.2.mmsa.output_projection.lora_A", "transformer.decoder_blocks.2.mmsa.output_projection.lora_B", "transformer.decoder_blocks.2.ffn.fc1.lora_A", "transformer.decoder_blocks.2.ffn.fc1.lora_B", "transformer.decoder_blocks.2.ffn.fc2.lora_A", "transformer.decoder_blocks.2.ffn.fc2.lora_B", "transformer.decoder_blocks.3.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.3.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.3.mmsa.output_projection.lora_A", "transformer.decoder_blocks.3.mmsa.output_projection.lora_B", "transformer.decoder_blocks.3.ffn.fc1.lora_A", "transformer.decoder_blocks.3.ffn.fc1.lora_B", "transformer.decoder_blocks.3.ffn.fc2.lora_A", "transformer.decoder_blocks.3.ffn.fc2.lora_B", "transformer.decoder_blocks.4.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.4.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.4.mmsa.output_projection.lora_A", "transformer.decoder_blocks.4.mmsa.output_projection.lora_B", "transformer.decoder_blocks.4.ffn.fc1.lora_A", "transformer.decoder_blocks.4.ffn.fc1.lora_B", "transformer.decoder_blocks.4.ffn.fc2.lora_A", "transformer.decoder_blocks.4.ffn.fc2.lora_B", "transformer.decoder_blocks.5.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.5.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.5.mmsa.output_projection.lora_A", "transformer.decoder_blocks.5.mmsa.output_projection.lora_B", "transformer.decoder_blocks.5.ffn.fc1.lora_A", "transformer.decoder_blocks.5.ffn.fc1.lora_B", "transformer.decoder_blocks.5.ffn.fc2.lora_A", "transformer.decoder_blocks.5.ffn.fc2.lora_B", "transformer.decoder_blocks.6.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.6.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.6.mmsa.output_projection.lora_A", "transformer.decoder_blocks.6.mmsa.output_projection.lora_B", "transformer.decoder_blocks.6.ffn.fc1.lora_A", "transformer.decoder_blocks.6.ffn.fc1.lora_B", "transformer.decoder_blocks.6.ffn.fc2.lora_A", "transformer.decoder_blocks.6.ffn.fc2.lora_B", "transformer.decoder_blocks.7.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.7.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.7.mmsa.output_projection.lora_A", "transformer.decoder_blocks.7.mmsa.output_projection.lora_B", "transformer.decoder_blocks.7.ffn.fc1.lora_A", "transformer.decoder_blocks.7.ffn.fc1.lora_B", "transformer.decoder_blocks.7.ffn.fc2.lora_A", "transformer.decoder_blocks.7.ffn.fc2.lora_B", "transformer.decoder_blocks.8.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.8.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.8.mmsa.output_projection.lora_A", "transformer.decoder_blocks.8.mmsa.output_projection.lora_B", "transformer.decoder_blocks.8.ffn.fc1.lora_A", "transformer.decoder_blocks.8.ffn.fc1.lora_B", "transformer.decoder_blocks.8.ffn.fc2.lora_A", "transformer.decoder_blocks.8.ffn.fc2.lora_B", "transformer.decoder_blocks.9.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.9.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.9.mmsa.output_projection.lora_A", "transformer.decoder_blocks.9.mmsa.output_projection.lora_B", "transformer.decoder_blocks.9.ffn.fc1.lora_A", "transformer.decoder_blocks.9.ffn.fc1.lora_B", "transformer.decoder_blocks.9.ffn.fc2.lora_A", "transformer.decoder_blocks.9.ffn.fc2.lora_B", "transformer.decoder_blocks.10.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.10.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.10.mmsa.output_projection.lora_A", "transformer.decoder_blocks.10.mmsa.output_projection.lora_B", "transformer.decoder_blocks.10.ffn.fc1.lora_A", "transformer.decoder_blocks.10.ffn.fc1.lora_B", "transformer.decoder_blocks.10.ffn.fc2.lora_A", "transformer.decoder_blocks.10.ffn.fc2.lora_B", "transformer.decoder_blocks.11.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.11.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.11.mmsa.output_projection.lora_A", "transformer.decoder_blocks.11.mmsa.output_projection.lora_B", "transformer.decoder_blocks.11.ffn.fc1.lora_A", "transformer.decoder_blocks.11.ffn.fc1.lora_B", "transformer.decoder_blocks.11.ffn.fc2.lora_A", "transformer.decoder_blocks.11.ffn.fc2.lora_B", "transformer.decoder_blocks.12.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.12.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.12.mmsa.output_projection.lora_A", "transformer.decoder_blocks.12.mmsa.output_projection.lora_B", "transformer.decoder_blocks.12.ffn.fc1.lora_A", "transformer.decoder_blocks.12.ffn.fc1.lora_B", "transformer.decoder_blocks.12.ffn.fc2.lora_A", "transformer.decoder_blocks.12.ffn.fc2.lora_B", "transformer.decoder_blocks.13.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.13.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.13.mmsa.output_projection.lora_A", "transformer.decoder_blocks.13.mmsa.output_projection.lora_B", "transformer.decoder_blocks.13.ffn.fc1.lora_A", "transformer.decoder_blocks.13.ffn.fc1.lora_B", "transformer.decoder_blocks.13.ffn.fc2.lora_A", "transformer.decoder_blocks.13.ffn.fc2.lora_B", "transformer.decoder_blocks.14.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.14.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.14.mmsa.output_projection.lora_A", "transformer.decoder_blocks.14.mmsa.output_projection.lora_B", "transformer.decoder_blocks.14.ffn.fc1.lora_A", "transformer.decoder_blocks.14.ffn.fc1.lora_B", "transformer.decoder_blocks.14.ffn.fc2.lora_A", "transformer.decoder_blocks.14.ffn.fc2.lora_B", "transformer.decoder_blocks.15.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.15.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.15.mmsa.output_projection.lora_A", "transformer.decoder_blocks.15.mmsa.output_projection.lora_B", "transformer.decoder_blocks.15.ffn.fc1.lora_A", "transformer.decoder_blocks.15.ffn.fc1.lora_B", "transformer.decoder_blocks.15.ffn.fc2.lora_A", "transformer.decoder_blocks.15.ffn.fc2.lora_B", "transformer.decoder_blocks.16.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.16.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.16.mmsa.output_projection.lora_A", "transformer.decoder_blocks.16.mmsa.output_projection.lora_B", "transformer.decoder_blocks.16.ffn.fc1.lora_A", "transformer.decoder_blocks.16.ffn.fc1.lora_B", "transformer.decoder_blocks.16.ffn.fc2.lora_A", "transformer.decoder_blocks.16.ffn.fc2.lora_B", "transformer.decoder_blocks.17.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.17.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.17.mmsa.output_projection.lora_A", "transformer.decoder_blocks.17.mmsa.output_projection.lora_B", "transformer.decoder_blocks.17.ffn.fc1.lora_A", "transformer.decoder_blocks.17.ffn.fc1.lora_B", "transformer.decoder_blocks.17.ffn.fc2.lora_A", "transformer.decoder_blocks.17.ffn.fc2.lora_B", "transformer.decoder_blocks.18.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.18.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.18.mmsa.output_projection.lora_A", "transformer.decoder_blocks.18.mmsa.output_projection.lora_B", "transformer.decoder_blocks.18.ffn.fc1.lora_A", "transformer.decoder_blocks.18.ffn.fc1.lora_B", "transformer.decoder_blocks.18.ffn.fc2.lora_A", "transformer.decoder_blocks.18.ffn.fc2.lora_B", "transformer.decoder_blocks.19.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.19.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.19.mmsa.output_projection.lora_A", "transformer.decoder_blocks.19.mmsa.output_projection.lora_B", "transformer.decoder_blocks.19.ffn.fc1.lora_A", "transformer.decoder_blocks.19.ffn.fc1.lora_B", "transformer.decoder_blocks.19.ffn.fc2.lora_A", "transformer.decoder_blocks.19.ffn.fc2.lora_B", "transformer.decoder_blocks.20.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.20.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.20.mmsa.output_projection.lora_A", "transformer.decoder_blocks.20.mmsa.output_projection.lora_B", "transformer.decoder_blocks.20.ffn.fc1.lora_A", "transformer.decoder_blocks.20.ffn.fc1.lora_B", "transformer.decoder_blocks.20.ffn.fc2.lora_A", "transformer.decoder_blocks.20.ffn.fc2.lora_B", "transformer.decoder_blocks.21.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.21.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.21.mmsa.output_projection.lora_A", "transformer.decoder_blocks.21.mmsa.output_projection.lora_B", "transformer.decoder_blocks.21.ffn.fc1.lora_A", "transformer.decoder_blocks.21.ffn.fc1.lora_B", "transformer.decoder_blocks.21.ffn.fc2.lora_A", "transformer.decoder_blocks.21.ffn.fc2.lora_B", "transformer.decoder_blocks.22.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.22.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.22.mmsa.output_projection.lora_A", "transformer.decoder_blocks.22.mmsa.output_projection.lora_B", "transformer.decoder_blocks.22.ffn.fc1.lora_A", "transformer.decoder_blocks.22.ffn.fc1.lora_B", "transformer.decoder_blocks.22.ffn.fc2.lora_A", "transformer.decoder_blocks.22.ffn.fc2.lora_B", "transformer.decoder_blocks.23.mmsa.qkv_projection.lora_A", "transformer.decoder_blocks.23.mmsa.qkv_projection.lora_B", "transformer.decoder_blocks.23.mmsa.output_projection.lora_A", "transformer.decoder_blocks.23.mmsa.output_projection.lora_B", "transformer.decoder_blocks.23.ffn.fc1.lora_A", "transformer.decoder_blocks.23.ffn.fc1.lora_B", "transformer.decoder_blocks.23.ffn.fc2.lora_A", "transformer.decoder_blocks.23.ffn.fc2.lora_B", "lm_head.lora_A", "lm_head.lora_B".

[T-AIMaven]

@javismiles Please check the reward model for matching state dicts. That's using GPT.from_checkpoint(), that's different from GPTRewardModel. When I have updated it into GPTRewardModel, this error is fixed.

@ethanyanjiali would you please check this together? Is it intended?

[abhishek-sharma-iisc]

@javismiles Please check the reward model for matching state dicts. That's using GPT.from_checkpoint(), that's different from GPTRewardModel. When I have updated it into GPTRewardModel, this error is fixed.

@ethanyanjiali would you please check this together? Is it intended?

@satel33 Can you please explain more clearly how this error need to be tackled. I'm also getting the same error.

[faizack]

I just fix the issue @javismiles @ethanyanjiali @satel33

Issue

The GPTRewardModel currently builds the value_head as a plain nn.Linear.
However, when cfg.lora_rank > 0, the checkpoint expects LoRA weights (lora_A, lora_B).
Because the class does not define a LoRA head, this leads to missing/unexpected key errors when loading LoRA-enabled checkpoints.

Fix

Update the GPTRewardModel class in gpt.py (line 489) to conditionally use a LoRA linear layer:

class GPTRewardModel(nn.Module):

    def __init__(self, cfg: TrainingConfig) -> None:
        super().__init__()
        self.cfg = cfg
        self.backbone = GPT(cfg)
        self.backbone.lm_head = nn.Identity()
        # Use LoRA for value_head if LoRA is enabled
        if cfg.lora_rank > 0:
            self.value_head = lora.Linear(
                cfg.embedding_dim, 1, bias=False, r=cfg.lora_rank
            )
        else:
            self.value_head = nn.Linear(cfg.embedding_dim, 1, bias=False)

    def forward(self, x: Tensor, attention_mask: Tensor = None):
        hidden = self.backbone(x, attention_mask)
        score = self.value_head(hidden).mean(dim=1)
        return score

    def freeze_weights(self, finetune_method):
        if finetune_method == "lora" and self.cfg.lora_rank > 0:
            lora.mark_only_lora_as_trainable(self)
        elif finetune_method == "last_block":
            trainable_params = [
                "backbone.transformer.decoder_blocks.35.mmsa.mask",
                "backbone.transformer.decoder_blocks.35.mmsa.qkv_projection.weight",
                "backbone.transformer.decoder_blocks.35.mmsa.qkv_projection.bias",
                "backbone.transformer.decoder_blocks.35.mmsa.output_projection.weight",
                "backbone.transformer.decoder_blocks.35.mmsa.output_projection.bias",
                "backbone.transformer.decoder_blocks.35.ln2.weight",
                "backbone.transformer.decoder_blocks.35.ln2.bias",
                "backbone.transformer.decoder_blocks.35.ffn.fc1.weight",
                "backbone.transformer.decoder_blocks.35.ffn.fc1.bias",
                "backbone.transformer.decoder_blocks.35.ffn.fc2.weight",
                "backbone.transformer.decoder_blocks.35.ffn.fc2.bias",
                "backbone.transformer.ln.weight",
                "backbone.transformer.ln.bias",
                "value_head.weight",
            ]
            for name, param in self.named_parameters():
                if name not in trainable_params:
                    param.requires_grad = False
                else:
                    print(f"{name} is trainable.")
        else:
            print(
                f"Unsupported method {finetune_method} (lora rank = {self.cfg.lora_rank})"
            )

    @classmethod
    def from_backbone_checkpoint(cls, cfg: TrainingConfig, ckpt_path: str):
        cfg.pretrain = ckpt_path
        model = GPTRewardModel(cfg)

        # Load the checkpoint to check if it has LoRA weights
        checkpoint = torch.load(ckpt_path, map_location="cpu")
        state_dict = checkpoint["model_state_dict"]

        # Check if the checkpoint contains LoRA weights
        has_lora = any("lora_A" in key or "lora_B" in key for key in state_dict.keys())

        if has_lora and cfg.lora_rank == 0:
            # If checkpoint has LoRA but current config doesn't, we need to adjust the config
            print(
                "Warning: Checkpoint contains LoRA weights but current config has lora_rank=0. Adjusting config to match checkpoint."
            )
            # We need to determine the LoRA rank from the checkpoint
            for key in state_dict.keys():
                if "lora_A" in key:
                    # Extract rank from the shape of lora_A weight
                    cfg.lora_rank = state_dict[key].shape[0]
                    break

        # Create the backbone with the adjusted config
        model.backbone = GPT(cfg)
        model.backbone.lm_head = nn.Identity()

        # Load state dict with strict=False to handle any remaining mismatches
        missing_keys, unexpected_keys = model.backbone.load_state_dict(
            state_dict, strict=False
        )

        if missing_keys:
            print(
                f"Missing keys when loading checkpoint: {missing_keys[:10]}..."
            )  # Show first 10
        if unexpected_keys:
            print(
                f"Unexpected keys when loading checkpoint: {unexpected_keys[:10]}..."
            )  # Show first 10

        return model

    @classmethod
    def from_checkpoint(
        cls, cfg: TrainingConfig, ckpt_path: str, strict=False, compile=False
    ):
        model = GPTRewardModel(cfg)
        if compile:
            model = torch.compile(model)
        checkpoint = torch.load(ckpt_path, map_location="cpu")
        model.load_state_dict(checkpoint["model_state_dict"], strict=strict)
        return model

    @classmethod
    def from_pretrained(cls, cfg: TrainingConfig):
        model = GPTRewardModel(cfg)
        model.backbone = GPT.from_pretrained(cfg)
        model.backbone.lm_head = nn.Identity()
        # model_states_keys = model.state_dict().keys()
        # with open('rm_states_keys.txt', 'w') as fp:
        #     for k in model_states_keys:
        #         fp.write(k + '\n')
        return model