The model training and model initialization code has been updated. (CPU/GPU auto-detection has been added.)

model_init.py

@@ -8,6 +8,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import sentencepiece as spm
+from config.model_config import *
 
 NUM_THREADS = os.cpu_count() or 4
 torch.set_num_threads(NUM_THREADS)
@@ -16,8 +17,17 @@ try:
 except RuntimeError:
     pass
 
-DEVICE = torch.device("cpu")
+# --- Automatic device selection: use a compatible GPU if available, otherwise fall back to CPU ---
+if torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    DEVICE = torch.device("mps")
+else:
+    DEVICE = torch.device("cpu")
+
 print(f"🧵 Number of CPU threads : {NUM_THREADS}")
+print(f"🖥️  Selected device      : {DEVICE.type.upper()}"
+      + (f" ({torch.cuda.get_device_name(0)})" if DEVICE.type == "cuda" else ""))
 
 class Vocab:
     PAD = "<pad>"
@@ -152,6 +162,10 @@ class OSW1Model(nn.Module):
     @torch.no_grad()
     def generate(self, idx, max_new_tokens, temperature=0.85, top_k=40, eos_id=None):
         self.eval()
+        # Make sure the input tensor lives on the same device as the model itself,
+        # so generation works no matter which device the checkpoint was trained on.
+        model_device = next(self.parameters()).device
+        idx = idx.to(model_device)
         for _ in range(max_new_tokens):
             idx_cond = idx[:, -self.block_size:]
             logits = self(idx_cond)
@@ -175,6 +189,8 @@ def find_checkpoint():
 
 def load_checkpoint(path: str):
     print(f"📦 Loading: {path}")
+    # map_location="cpu" guarantees the checkpoint can always be read back,
+    # regardless of which device (GPU/MPS/CPU) it was trained on.
     ckpt = torch.load(path, map_location="cpu")
 
     cfg = ckpt["config"]
@@ -188,6 +204,7 @@ def load_checkpoint(path: str):
     param_count = ckpt.get("param_count", sum(p.numel() for p in model.parameters()))
     training_time = ckpt.get("training_time_sec", None)
     final_loss = ckpt.get("final_loss", None)
+    trained_on = ckpt.get("trained_on_device", "unknown")
 
     print("\n" + "=" * 64)
     print("🧠  OpenSoftware-World OSW1 — LOADED MODEL INFORMATION")
@@ -202,6 +219,7 @@ def load_checkpoint(path: str):
         print(f"  Training time        : {training_time/60:.2f} minutes")
     if final_loss is not None:
         print(f"  Final training loss  : {final_loss:.4f}")
+    print(f"  Trained on device    : {trained_on}  ->  Running on: {DEVICE.type}")
     print("=" * 64 + "\n")
 
     return model, vocab, cfg
@@ -211,6 +229,7 @@ def chat_loop(model: OSW1Model, vocab: Vocab):
     print("💬 OpenSoftware-World-OSW1 ready! You can start chatting. Type 'exit' to quit.")
     print("=" * 64)
 
+    model_device = next(model.parameters()).device
     eos_id = vocab.stoi[Vocab.EOS]
     bos_id = vocab.stoi[Vocab.BOS]
 
@@ -228,8 +247,8 @@ def chat_loop(model: OSW1Model, vocab: Vocab):
             continue
 
         ids = [bos_id] + vocab.encode(user_in)
-        x = torch.tensor([ids], dtype=torch.long)
-        out = model.generate(x, max_new_tokens=60, temperature=0.85, top_k=40, eos_id=eos_id)
+        x = torch.tensor([ids], dtype=torch.long, device=model_device)
+        out = model.generate(x, max_new_tokens=init_max_new_tokens, temperature=init_temperature, top_k=init_top_k, eos_id=eos_id)
         answer_ids = out[0, len(ids):].tolist()
         answer = vocab.decode(answer_ids)
         print(f"OpenSoftware-World-OSW1: {answer if answer else '(...silence...)'}")
@@ -255,4 +274,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()

model_training.py

@@ -5,12 +5,15 @@ import math
 import time
 import glob
 import random
+import contextlib
 
 from dataclasses import dataclass
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import sentencepiece as spm
+from config.model_config import *
 
 torch.manual_seed(42)
 random.seed(42)
@@ -28,84 +31,111 @@ try:
 except Exception:
     pass
 
-DEVICE = torch.device("cpu")
+# --- Automatic device selection: use a compatible GPU if available, otherwise fall back to CPU ---
+if torch.cuda.is_available():
+    DEVICE = torch.device("cuda")
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    DEVICE = torch.device("mps")
+else:
+    DEVICE = torch.device("cpu")
 
 # Autocasting to bfloat16 on the CPU can speed up most matmul operations (if supported)
-USE_BF16_AUTOCAST = True
-try:
-    _ = torch.zeros(1, dtype=torch.bfloat16) + torch.zeros(1, dtype=torch.bfloat16)
-except Exception:
-    USE_BF16_AUTOCAST = False
+# This optimization only matters when we are actually training on CPU.
+USE_BF16_AUTOCAST = False
+if DEVICE.type == "cpu":
+    try:
+        _ = torch.zeros(1, dtype=torch.bfloat16) + torch.zeros(1, dtype=torch.bfloat16)
+        USE_BF16_AUTOCAST = True
+    except Exception:
+        USE_BF16_AUTOCAST = False
+
+# Whether bfloat16 autocast is available on the current CUDA GPU (Ampere+ generally supports this)
+USE_CUDA_BF16_AUTOCAST = DEVICE.type == "cuda" and torch.cuda.is_bf16_supported()
 
 print(f"🧵 Number of CPU threads           : {NUM_THREADS}")
-print(f"⚙️  bfloat16 autocast status    : {'active' if USE_BF16_AUTOCAST else 'inactive'}")
+print(f"🖥️  Selected training device    : {DEVICE.type.upper()}"
+      + (f" ({torch.cuda.get_device_name(0)})" if DEVICE.type == "cuda" else ""))
+print(f"⚙️  bfloat16 autocast status    : "
+      f"{'active (CPU)' if USE_BF16_AUTOCAST else ('active (CUDA)' if USE_CUDA_BF16_AUTOCAST else 'inactive')}")
 
 @dataclass
 class OSW1Config:
     data_dir: str = "data"
 
-    block_size: int = 128      
-    d_model: int = 256         
-    n_layer: int = 4           
-    n_head: int = 4            
-    d_ff: int = 1024           
-    dropout: float = 0.1
-
-    batch_size: int = 8
-    grad_accum_steps: int = 2      
-    epochs: int = 20
-    max_lr: float = 3e-4
-    min_lr: float = 3e-5
-    warmup_ratio: float = 0.05
-    weight_decay: float = 0.1
-    grad_clip: float = 1.0
-    label_smoothing: float = 0.05
+    block_size: int = training_block_size     
+    d_model: int = training_d_model         
+    n_layer: int = training_n_layer           
+    n_head: int = training_n_head            
+    d_ff: int = training_d_ff           
+    dropout: float = training_dropout
+
+    batch_size: int = training_batch_size
+    grad_accum_steps: int = training_grad_accum_steps      
+    epochs: int = training_epochs
+    max_lr: float = training_max_lr
+    min_lr: float = training_min_lr
+    warmup_ratio: float = training_warmup_ratio
+    weight_decay: float = training_weight_decay
+    grad_clip: float = training_grad_clip
+    label_smoothing: float = training_label_smoothing
 
     checkpoint_prefix: str = "opensoftware_world_osw1"
 
-TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)
-
-def tokenize(text: str):
-    return TOKEN_RE.findall(text.lower())
-
 class Vocab:
-    PAD, UNK, BOS, EOS = "<pad>", "<unk>", "<bos>", "<eos>"
+    PAD = "<pad>"
+    UNK = "<unk>"
+    BOS = "<bos>"
+    EOS = "<eos>"
 
-    def __init__(self):
-        self.stoi = {}
-        self.itos = []
+    def __init__(self, model_path="opensoftware_world_osw1_tokenizer.model"):
+        self.sp = spm.SentencePieceProcessor()
+        self.sp.load(model_path)
 
-    def build(self, token_stream):
-        specials = [Vocab.PAD, Vocab.UNK, Vocab.BOS, Vocab.EOS]
-        counts = {}
-        for tok in token_stream:
-            counts[tok] = counts.get(tok, 0) + 1
-        sorted_toks = sorted(counts.items(), key=lambda x: (-x[1], x[0]))
-        self.itos = specials + [t for t, _ in sorted_toks]
-        self.stoi = {t: i for i, t in enumerate(self.itos)}
+        self.pad_id = self.sp.pad_id()
+        self.unk_id = self.sp.unk_id()
+        self.bos_id = self.sp.bos_id()
+        self.eos_id = self.sp.eos_id()
 
     def encode(self, text, add_bos=False, add_eos=False):
-        ids = [self.stoi.get(t, self.stoi[Vocab.UNK]) for t in tokenize(text)]
+        ids = self.sp.encode(text, out_type=int)
+
         if add_bos:
-            ids = [self.stoi[Vocab.BOS]] + ids
+            ids = [self.bos_id] + ids
+
         if add_eos:
-            ids = ids + [self.stoi[Vocab.EOS]]
+            ids = ids + [self.eos_id]
+
         return ids
 
     def decode(self, ids):
-        toks = [self.itos[i] for i in ids if 0 <= i < len(self.itos)]
-        toks = [t for t in toks if t != Vocab.PAD and t != Vocab.BOS]
-        out = []
-        for t in toks:
-            if t == Vocab.EOS:
-                break
-            out.append(t)
-        text = " ".join(out)
-        text = re.sub(r"\s+([.,!?;:])", r"\1", text)
-        return text
+        ids = [
+            i for i in ids
+            if i not in (self.pad_id, self.bos_id)
+        ]
+
+        if self.eos_id in ids:
+            ids = ids[:ids.index(self.eos_id)]
+
+        return self.sp.decode(ids)
 
     def __len__(self):
-        return len(self.itos)
+        return self.sp.get_piece_size()
+
+    @property
+    def stoi(self):
+        return {
+            self.PAD: self.pad_id,
+            self.UNK: self.unk_id,
+            self.BOS: self.bos_id,
+            self.EOS: self.eos_id,
+        }
+
+    @property
+    def itos(self):
+        return [
+            self.sp.id_to_piece(i)
+            for i in range(self.sp.get_piece_size())
+        ]
 
 def load_json_pairs(json_dir):
     pairs = []
@@ -179,14 +209,6 @@ def build_corpus(cfg: OSW1Config, vocab: Vocab):
             "folders with data for the model to learn from."
         )
 
-    all_tokens = []
-    for q, a in qa_pairs:
-        all_tokens.extend(tokenize(q))
-        all_tokens.extend(tokenize(a))
-    for t in plain_texts:
-        all_tokens.extend(tokenize(t))
-    vocab.build(all_tokens)
-
     sequences = []
 
     for q, a in qa_pairs:
@@ -333,6 +355,10 @@ class OSW1Model(nn.Module):
     def generate(self, idx, max_new_tokens, temperature=0.9, top_k=40, eos_id=None):
         was_training = self.training
         self.eval()
+        # Make sure the input tensor lives on the same device as the model itself,
+        # so generation works no matter which device the model was trained/loaded on.
+        model_device = next(self.parameters()).device
+        idx = idx.to(model_device)
         for _ in range(max_new_tokens):
             idx_cond = idx[:, -self.cfg.block_size:]
             logits, _ = self(idx_cond)
@@ -402,6 +428,20 @@ def lr_at_step(step, total_steps, warmup_steps, max_lr, min_lr):
     progress = min(max(progress, 0.0), 1.0)
     return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * progress))
 
+def get_autocast_context():
+    """
+    Returns the correct autocast context manager for whichever device we ended up
+    training on (CUDA, CPU, or MPS/other). Falls back to a no-op context if the
+    current device doesn't support (or benefit from) autocasting here.
+    """
+    if DEVICE.type == "cuda":
+        dtype = torch.bfloat16 if USE_CUDA_BF16_AUTOCAST else torch.float16
+        return torch.autocast(device_type="cuda", dtype=dtype)
+    elif DEVICE.type == "cpu" and USE_BF16_AUTOCAST:
+        return torch.autocast(device_type="cpu", dtype=torch.bfloat16)
+    else:
+        return contextlib.nullcontext()
+
 def train(cfg: OSW1Config):
     vocab = Vocab()
     sequences = build_corpus(cfg, vocab)
@@ -440,6 +480,12 @@ def train(cfg: OSW1Config):
         weight_decay=cfg.weight_decay,
     )
 
+    # Only needed for numerically-fragile float16 training on CUDA GPUs that lack
+    # native bfloat16 support. When bfloat16 is available (or we're on CPU/MPS),
+    # the scaler simply stays disabled and behaves as a no-op.
+    use_grad_scaler = DEVICE.type == "cuda" and not USE_CUDA_BF16_AUTOCAST
+    scaler = torch.amp.GradScaler(enabled=use_grad_scaler)
+
     steps_per_epoch = max(1, len(loader) // cfg.grad_accum_steps)
     total_steps = steps_per_epoch * cfg.epochs
     warmup_steps = max(1, int(total_steps * cfg.warmup_ratio))
@@ -459,21 +505,21 @@ def train(cfg: OSW1Config):
         for i, (x, y) in enumerate(loader):
             x, y = x.to(DEVICE), y.to(DEVICE)
 
-            if USE_BF16_AUTOCAST:
-                with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
-                    _, loss = compiled_model(x, y)
-            else:
+            with get_autocast_context():
                 _, loss = compiled_model(x, y)
 
             loss_scaled = loss / cfg.grad_accum_steps
-            loss_scaled.backward()
+            scaler.scale(loss_scaled).backward()
 
             if (i + 1) % cfg.grad_accum_steps == 0:
+                if use_grad_scaler:
+                    scaler.unscale_(optimizer)
                 torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
                 lr = lr_at_step(global_step, total_steps, warmup_steps, cfg.max_lr, cfg.min_lr)
                 for g in optimizer.param_groups:
                     g["lr"] = lr
-                optimizer.step()
+                scaler.step(optimizer)
+                scaler.update()
                 optimizer.zero_grad(set_to_none=True)
                 global_step += 1
 
@@ -497,15 +543,18 @@ def train(cfg: OSW1Config):
           f"{total_time/60:.2f} minutes ({total_time:.1f} seconds)\n")
 
     ckpt_path = f"{cfg.checkpoint_prefix}_{size_tag}.pth"
+    # Move every tensor in the state dict to CPU before saving. This makes the
+    # checkpoint device-agnostic: a model trained on GPU can later be loaded
+    # and run correctly on a machine that only has a CPU (and vice versa).
+    cpu_state_dict = {k: v.detach().cpu() for k, v in model.state_dict().items()}
     torch.save({
-        "model_state_dict": model.state_dict(),
+        "model_state_dict": cpu_state_dict,
         "config": cfg.__dict__,
-        "vocab_stoi": vocab.stoi,
-        "vocab_itos": vocab.itos,
         "pad_id": pad_id,
         "param_count": sum(p.numel() for p in model.parameters()),
         "training_time_sec": total_time,
         "final_loss": avg_loss,
+        "trained_on_device": DEVICE.type,
     }, ckpt_path)
     print(f"💾 Model saved: {ckpt_path}\n")
 
@@ -516,12 +565,13 @@ def chat_loop(model: OSW1Model, vocab: Vocab, cfg: OSW1Config):
     print("💬 OSW1 with chat mode! Type 'exit' to quit.")
     print("=" * 64)
     model.eval()
+    model_device = next(model.parameters()).device
     eos_id = vocab.stoi[Vocab.EOS]
     bos_id = vocab.stoi[Vocab.BOS]
 
     while True:
         try:
-            user_in = input("\You: ").strip()
+            user_in = input("\nYou: ").strip()
         except (EOFError, KeyboardInterrupt):
             print("\n👋 Goodbye!")
             break
@@ -533,21 +583,28 @@ def chat_loop(model: OSW1Model, vocab: Vocab, cfg: OSW1Config):
             continue
 
         ids = [bos_id] + vocab.encode(user_in)
-        x = torch.tensor([ids], dtype=torch.long)
-        out = model.generate(x, max_new_tokens=60, temperature=0.85, top_k=40, eos_id=eos_id)
+        x = torch.tensor([ids], dtype=torch.long, device=model_device)
+        out = model.generate(x, max_new_tokens=training_max_new_tokens, temperature=training_temperature, top_k=training_top_k, eos_id=eos_id)
         answer_ids = out[0, len(ids):].tolist()
         answer = vocab.decode(answer_ids)
         print(f"OSW1: {answer if answer else '(...silence...)'}")
 
 def load_checkpoint(path: str):
+    # map_location="cpu" guarantees the checkpoint can always be read back,
+    # regardless of which device it was trained on or whether a GPU is present
+    # on the machine doing the loading.
     ckpt = torch.load(path, map_location="cpu")
     cfg = OSW1Config(**ckpt["config"])
     vocab = Vocab()
-    vocab.stoi = ckpt["vocab_stoi"]
-    vocab.itos = ckpt["vocab_itos"]
     model = OSW1Model(len(vocab), cfg, pad_id=ckpt["pad_id"])
     model.load_state_dict(ckpt["model_state_dict"])
+    # Now move the freshly-loaded model onto whichever device is available
+    # on *this* machine (GPU/MPS if present, otherwise CPU) so it runs correctly
+    # no matter what device it was originally trained on.
+    model.to(DEVICE)
     model.eval()
+    trained_on = ckpt.get("trained_on_device", "unknown")
+    print(f"📦 Checkpoint loaded (trained on: {trained_on}) -> running on: {DEVICE.type}")
     return model, vocab, cfg
 
 def main():
@@ -557,4 +614,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()