From e8f15b56ea395921bfe932e0bf1be9db78e546d0 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Sat, 10 Oct 2026 01:40:39 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20BiLSTM=20trains=20in=20fp32=20with=20NaN?= =?UTF-8?q?=20guard=20=E2=80=94=20bf16=20autocast=20NaN-collapsed=20the=20?= =?UTF-8?q?LSTM=20(all-NaN=20checkpoint)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- train_bilstm_hf.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/train_bilstm_hf.py b/train_bilstm_hf.py index 2662cb9..fd2574f 100644 --- a/train_bilstm_hf.py +++ b/train_bilstm_hf.py @@ -159,9 +159,12 @@ def forward(self, ids): for _ep in range(EPOCHS): for ids, lab in loader: ids, lab = ids.cuda(), lab.cuda() - with torch.autocast("cuda", torch.bfloat16): - logits = model(ids) + logits = model(ids) # fp32: LSTMs NaN-collapse under bf16 autocast loss = ce(logits.reshape(-1, n_classes), lab.reshape(-1)) + if not torch.isfinite(loss): + print(f"[nan] step {step}: non-finite loss, skipping batch", flush=True) + opt.zero_grad() + continue loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() @@ -186,8 +189,7 @@ def tag(text: str) -> str: skel, _ = HP.split_planes(text) ids = [vocab.get(ch, 1) for ch in skel] or [0] t = torch.tensor([ids], device="cuda") - with torch.autocast("cuda", torch.bfloat16): - logits = model(t) + logits = model(t) preds = logits[0].argmax(-1).tolist() return "".join(ch + (classes[preds[i]] if i < len(preds) else "") for i, ch in enumerate(skel))