tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2025-12-06 04:12:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 16788eed3c into 4a87a0d19f

Browse Source
This commit is contained in:
Sitananda Prasad 2025-11-23 20:14:26 +05:30 committed by GitHub
commit 6398858ea9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
nanochat/gpt.py
View File

@ -265,13 +265,14 @@ class GPT(nn.Module):
# training mode: compute and return the loss # training mode: compute and return the loss
# TODO: experiment with Liger Kernels / chunked cross-entropy etc. # TODO: experiment with Liger Kernels / chunked cross-entropy etc.
logits = self.lm_head(x) logits = self.lm_head(x)
logits = softcap * torch.tanh(logits / softcap) # logits softcap
logits = logits.float() # use tf32/fp32 for logits logits = logits.float() # use tf32/fp32 for logits
logits = softcap * torch.tanh(logits / softcap) # logits softcap
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction)
return loss return loss
else: else:
# inference mode: compute and return the logits # inference mode: compute and return the logits
logits = self.lm_head(x) logits = self.lm_head(x)
logits = logits.float() # use tf32/fp32 for logits
logits = softcap * torch.tanh(logits / softcap) # logits softcap logits = softcap * torch.tanh(logits / softcap) # logits softcap
return logits return logits