wuyang/llm-atlas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: flatten noncontiguous AttnRes targets

Browse Source
This commit is contained in:
wuyang
2026-07-30 06:39:29 +08:00
parent 1f20f81939
commit 5f49906be6
1 changed files with 5 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
experiments/k3/attnres/train.py
+5 -3
View File
@@ -390,7 +390,9 @@ def learning_rate(step: int, total_steps: int) -> float:
def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor: def cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
return F.cross_entropy(logits.float().view(-1, VOCABULARY), targets.view(-1)) return F.cross_entropy(
logits.float().reshape(-1, VOCABULARY), targets.reshape(-1)
)
@torch.no_grad() @torch.no_grad()
@@ -410,8 +412,8 @@ def evaluate(
with torch.autocast(device_type="cuda", dtype=torch.bfloat16): with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
logits, _ = model(inputs) logits, _ = model(inputs)
loss = F.cross_entropy( loss = F.cross_entropy(
logits.float().view(-1, VOCABULARY), logits.float().reshape(-1, VOCABULARY),
targets.view(-1), targets.reshape(-1),
reduction="sum", reduction="sum",
) )
loss_sum += loss.detach().cpu().item() loss_sum += loss.detach().cpu().item()