opt reward model

f7516c41 · u010280923 · 5b3f373a · f7516c41
显示空白变更内容
内联并排

Showing with 3 addition and 4 deletion

src/rlhf/reward.py src/rlhf/reward.py +3 -4

未找到文件。
--- a/src/rlhf/reward.py
+++ b/src/rlhf/reward.py
@@ -126,11 +126,10 @@ class RewardModel(pl.LightningModule):
                {"params": [p for n, p in self.named_parameters()], "weight_decay": 0.0},
            ]
        
-        # if self.deepspeed_offload:
-        #     return DeepSpeedCPUAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adamw_mode=False, weight_decay=0, amsgrad=False)
-        # return FusedAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adam_w_mode=False, weight_decay=0, amsgrad=False)
+        if self.deepspeed_offload:
+            return DeepSpeedCPUAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adamw_mode=False, weight_decay=0, amsgrad=False)
+        return FusedAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, adam_w_mode=False, weight_decay=0, amsgrad=False)
        # return ZeroOneAdam(optim_groups, lr=self.args.lr_init, betas=self.args.betas, eps=self.args.adam_eps, bias_correction=True, weight_decay=0, amsgrad=False, cuda_aware=False)
-        return torch.optim.Adam(optim_groups, lr=1e-5, betas=(0.9, 0.95))
    
    @property
    def deepspeed_offload(self) -> bool: