Update grok 1 model (#1095)

This commit is contained in:
Lianmin Zheng
2024-08-14 04:40:44 -07:00
committed by GitHub
parent fe5024325b
commit a59636bb5e
11 changed files with 814 additions and 514 deletions

View File

@@ -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