Device: use CUDA when available; model.to(device) after loading works

README.md

@@ -75,8 +75,9 @@ Notes:
 - The checkpoint is not a standard architecture, so it has no `forward` for general use: a poem is produced by the
   project's own decoder (`src/` in this repository, unchanged from the paper's code), which the wrapper calls.
 - The first load also downloads the WangchanBERTa encoder the model is built on.
-- Device: Apple GPU (mps) when available, otherwise CPU; pass `device="cpu"` to `from_pretrained` to choose. CUDA has
-  not been tested.
+- Device: an NVIDIA GPU (cuda) when available, then the Apple GPU (mps), otherwise CPU; pass `device="cuda"`,
+  `"mps"` or `"cpu"` to `from_pretrained` to choose, or call `model.to(device)` after loading. Tested on mps and CPU;
+  CUDA has not been tested. The decoder's masks and sampling run on the CPU, so a GPU speeds up only the network pass.
 - Decoding is deterministic given the weights and the seed, and the wrapper returns exactly the stanzas of the paper's
   decode. A stanza takes about 1.6 seconds on a laptop GPU.
 - `src/` holds the generation code and the model and training code (`src/scripts/twolayer_canvas.py`). The data-building and

modeling_klon8.py

@@ -11,8 +11,8 @@ loads the checkpoint with the project's code, shipped unchanged in `src/` of thi
 `generate_poem`. With the same seed it returns exactly the stanzas of the paper's decode.
 
 Requirements: torch, transformers, sentencepiece, pythainlp, tltk. The first load also downloads the WangchanBERTa
-encoder (`config.base_model`), which the model is built on. Device: Apple GPU (mps) when available, else CPU; pass
-`device="cpu"` (or set KLON8_DEVICE) to choose. One klon8 checkpoint can be loaded per Python process.
+encoder (`config.base_model`), which the model is built on. Device: NVIDIA GPU (cuda) when available, then Apple GPU (mps), else CPU;
+pass `device=...` (or set KLON8_DEVICE) to choose, or call `model.to(device)` after loading. One klon8 checkpoint can be loaded per Python process.
 """
 import importlib
 import os

src/scripts/twolayer_canvas.py

@@ -28,7 +28,7 @@ import torch, torch.nn as nn, torch.nn.functional as F
 from klon8 import prosody as P
 
 WC = "airesearch/wangchanberta-base-att-spm-uncased"
-DEV = torch.device(os.environ.get("KLON8_DEVICE") or ("mps" if torch.backends.mps.is_available() else "cpu"))
+DEV = torch.device(os.environ.get("KLON8_DEVICE") or ("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"))
 C, TL, L = 10, 12, 4
 T = TL + L * C
 PAD, MASK, UNK, SEP = 0, 1, 2, 3                         # word ids AND token ids (specials have k=0, n=1)
@@ -547,7 +547,8 @@ class Decoder:
         """masked log-probs [T, NT] for the current state"""
         X, tk, tn, tw = self.X, self.X.tok_k, self.X.tok_n, self.X.tok_word
         self.n_fwd += 1
-        logits, _ = self.m(self.tok[None].to(DEV), self.snd[None].to(DEV), self.ri_in[None].to(DEV))
+        dev = self.m.word_bias.device                     # the model's own device, so model.to(...) after loading works
+        logits, _ = self.m(self.tok[None].to(dev), self.snd[None].to(dev), self.ri_in[None].to(dev))
         lg = logits[0].float().cpu() + self.base[None, :]
         for i in range(L):
             lg[self.idx(i, 0)] += torch.where(tk == 0, 0.0, -1e9)