Update grok 1 model (#1095)
This commit is contained in:
@@ -164,9 +164,9 @@ class LogitsProcessor(nn.Module):
|
||||
last_logits = last_logits[:, : self.config.vocab_size].float()
|
||||
|
||||
if hasattr(self.config, "final_logit_softcapping"):
|
||||
last_logits /= self.config.final_logit_softcapping
|
||||
last_logits.div_(self.config.final_logit_softcapping)
|
||||
last_logits = torch.tanh(last_logits)
|
||||
last_logits *= self.config.final_logit_softcapping
|
||||
last_logits.mul_(self.config.final_logit_softcapping)
|
||||
|
||||
# Return only last_logits if logprob is not requested
|
||||
if not logits_metadata.return_logprob:
|
||||
@@ -209,9 +209,9 @@ class LogitsProcessor(nn.Module):
|
||||
all_logits = all_logits[:, : self.config.vocab_size].float()
|
||||
|
||||
if hasattr(self.config, "final_logit_softcapping"):
|
||||
all_logits /= self.config.final_logit_softcapping
|
||||
all_logits.div_(self.config.final_logit_softcapping)
|
||||
all_logits = torch.tanh(all_logits)
|
||||
all_logits *= self.config.final_logit_softcapping
|
||||
all_logits.mul_(self.config.final_logit_softcapping)
|
||||
|
||||
all_logprobs = all_logits
|
||||
del all_logits, hidden_states
|
||||
|
||||
Reference in New Issue
Block a user