| 128 | dist.destroy_process_group() |
| 129 | |
| 130 | class YieldPredLayer(nn.Module): |
| 131 | def __init__(self, input_size, hidden_size, output_size=1): |
| 132 | super(YieldPredLayer, self).__init__() |
| 133 | self.act = nn.SiLU() |
| 134 | self.predictor = nn.Sequential( |
| 135 | nn.Linear(input_size, hidden_size), |
| 136 | # self.act, |
| 137 | # nn.Linear(hidden_size, hidden_size//4), |
| 138 | # # self.act, |
| 139 | nn.Linear(hidden_size, 1), |
| 140 | ) |
| 141 | def forward(self, x): |
| 142 | pred = self.predictor(x) |
| 143 | # print(f'pred:{pred.view(-1)}') |
| 144 | # print(f'y: {y.view(-1)}') |
| 145 | return pred |
| 146 | |
| 147 | class LlamaWithLoss(nn.Module): |
| 148 | def __init__(self, llama, predictor): |