gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

model: silence torch parity warning (read loss before backward)

Browse Source 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-15 16:09:30 +08:00
parent 3366f30c4d
commit 603c85e1e0
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
crates/xtrain-model/tests/parity.py
View File

@@ -134,6 +134,7 @@ logits = h @ lm_head # [seq, vocab]
loss = torch.nn.functional.cross_entropy( loss = torch.nn.functional.cross_entropy(
logits, torch.tensor(targets, dtype=torch.long), reduction="mean") logits, torch.tensor(targets, dtype=torch.long), reduction="mean")
loss_val = loss.item()
loss.backward() loss.backward()
# ---- Compare ---- # ---- Compare ----
@@ -147,8 +148,8 @@ def relerr(a, b):
ref_logits = read_vec("logits.txt") ref_logits = read_vec("logits.txt")
ref_loss = read_vec("loss.txt").item() ref_loss = read_vec("loss.txt").item()
print(f"loss: rust={ref_loss:.6e} torch={loss.item():.6e} " print(f"loss: rust={ref_loss:.6e} torch={loss_val:.6e} "
f"relerr={abs(loss.item()-ref_loss)/max(abs(ref_loss),1e-6):.2e}") f"relerr={abs(loss_val-ref_loss)/max(abs(ref_loss),1e-6):.2e}")
le = relerr(logits.detach(), ref_logits) le = relerr(logits.detach(), ref_logits)
print(f"logits: max relerr = {le:.2e}") print(f"logits: max relerr = {le:.2e}")