Update fine_tune_jit_with_validation_1b.py

fine_tune_jit_with_validation_1b.py

@@ -1,340 +1,356 @@
-# Copyright (c) 2025 CMS Manhattan
-# All rights reserved.
-#
-# This file is part of a project authored by CMS Manhattan. You may use, distribute, and modify
-# this code under the terms of the APACHE 2.0 license.
-
-
-"""
-Before run this script, download the GPT-2 tokenizer files into a local folder named 'tokenizer':
-mkdir -p tokenizer
-wget -O tokenizer/tokenizer.json          https://huggingface.co/gpt2/resolve/main/tokenizer.json
-wget -O tokenizer/vocab.json              https://huggingface.co/gpt2/resolve/main/vocab.json
-wget -O tokenizer/merges.txt              https://huggingface.co/gpt2/resolve/main/merges.txt
-wget -O tokenizer/tokenizer_config.json   https://huggingface.co/gpt2/resolve/main/tokenizer_config.json
-"""
-import os
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import IterableDataset, DataLoader
-from transformers import GPT2TokenizerFast
-from tqdm import tqdm
-import shutil
-import math
-from pathlib import Path
-import re
-
-# ============================= SETTINGS =============================
-TRAIN_SEQ_LEN = 256
-BATCH_SIZE = 1
-EPOCHS = 1
-LEARNING_RATE = 6e-6
-WEIGHT_DECAY = 0.01
-GRAD_CLIP = 1.0
-KEEP_LAST_EPOCHS = 3
-VAL_SPLIT_RATIO = 0.05
-VOCAB_SIZE = 50257
-
-BASE_MODEL_PATH = Path("models/gpt_modern_1b_class.script.pt")
-LAST_TRAINED_PATH = Path("models/gpt_1b_last_trained.script.pt")
-BACKUP_DIR = Path("models/backups")
-BACKUP_DIR.mkdir(exist_ok=True)
-
-RAW_PATH = Path("datasets/dialogues_text.txt")
-CLEAN_PATH = Path("datasets/dialogues_text_clean.txt")
-
-# Device selection
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
-
-# -- Dataset cleaning --
-force_clean = False
-if not CLEAN_PATH.exists():
-    print("Cleaned dataset not found. Performing initial cleaning...")
-    force_clean = True
-else:
-    try:
-        if RAW_PATH.stat().st_mtime > CLEAN_PATH.stat().st_mtime:
-            print("Detected changes in the raw dataset. Re-cleaning...")
-            force_clean = True
-        else:
-            print(f"Using existing cleaned dataset → {CLEAN_PATH}")
-    except FileNotFoundError:
-        print("File system synchronization error. Performing re-cleaning for safety...")
-        force_clean = True
-
-if force_clean:
-    if not RAW_PATH.exists():
-        raise FileNotFoundError(f"ERROR: Source file {RAW_PATH} not found. Check the path.")
-
-    print("Cleaning up the dataset from garbage (wrong separators, extra spaces)...")
-    text = RAW_PATH.read_text(encoding="utf-8")
-    text = re.sub(r' {2,}', ' ', text)
-    text = text.replace(" \n", "\n").replace("\n ", "\n")
-    CLEAN_PATH.write_text(text, encoding="utf-8")
-    print(f"Dataset successfully cleaned and saved → {CLEAN_PATH}")
-
-DATASET_PATH = CLEAN_PATH
-OUTPUT_DIR = Path("build/fine_tuning_output")
-MODEL_SAVE_NAME = "gpt_finetuned.script.pt"
-
-# ============================= DATASET (LAZY) =============================
-
-class LazyTextDataset(IterableDataset):
-    """Lazy memory-efficient dataset, splits on-the-fly into train and val."""
-    # Обратите внимание: аргумент tokenizer_name по-прежнему имеет значение по умолчанию "gpt2",
-    # но в функции train() мы теперь передаем локальный путь.
-    def __init__(self, text_file, seq_len=TRAIN_SEQ_LEN, tokenizer_name="gpt2", split_type='train', val_ratio=VAL_SPLIT_RATIO):
-        self.seq_len = seq_len
-        # Эта строка теперь загружает токенизатор из локальной папки, если передан локальный путь.
-        self.tokenizer = GPT2TokenizerFast.from_pretrained(tokenizer_name)
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-        self.text_file = text_file
-        self.split_type = split_type
-        self.val_ratio = val_ratio
-
-        print(f"Loading and tokenizing text from {text_file}")
-        with open(text_file, "r", encoding="utf-8") as f:
-            self.data = f.read()
-        self.tokens = self.tokenizer.encode(self.data)
-        
-        # Work out split indices
-        total_tokens = len(self.tokens) - 1  # because label sequence shifted
-        total_batches = total_tokens // seq_len
-        val_size = int(total_batches * self.val_ratio)
-        train_size = total_batches - val_size
-        if split_type == 'train':
-            self.start = 0
-            self.stop = train_size
-        elif split_type == 'val':
-            self.start = train_size
-            self.stop = train_size + val_size
-        else:
-            raise ValueError(f"split_type should be 'train' or 'val', got {split_type}")
-        self.total_sequences = self.stop - self.start
-        print(f"Lazy dataset: {self.total_sequences:,} sequences for {split_type} split (from {total_batches:,} total)")
-
-    def __iter__(self):
-        for i in range(self.start * self.seq_len, self.stop * self.seq_len, self.seq_len):
-            # Make sure last batch fits
-            if i + self.seq_len + 1 > len(self.tokens):
-                break
-            input_seq = torch.tensor(self.tokens[i : i + self.seq_len], dtype=torch.long)
-            label_seq = torch.tensor(self.tokens[i + 1 : i + self.seq_len + 1], dtype=torch.long)
-            yield input_seq, label_seq
-
-    def __len__(self):
-        return self.total_sequences
-
-# ============================= GET LOGITS UTIL =============================
-
-def get_logits_from_model(model, inputs):
-    """
-    Robust wrapper to call either a scripted JIT model or nn.Module.
-    Handles models that either return (logits, kv) or just logits.
-    """
-    # Ensure inputs on same device as model parameters/buffers
-    inputs = inputs.to(device)
-    try:
-        out = model(inputs)
-    except RuntimeError as e:
-        # Some JIT modules expect plain tensor on CPU device for tracing path.
-        # Re-raise if unrelated
-        raise
-
-    # Model may return logits or (logits, kv)
-    if isinstance(out, tuple) or (isinstance(out, list) and len(out) >= 1):
-        logits = out[0]
-    else:
-        logits = out
-    return logits
-
-# ============================= EVALUATION (VALIDATION) =============================
-
-def evaluate(model, dataloader, criterion, device):
-    model.eval()
-    total_loss = 0.0
-    count = 0
-    with torch.no_grad():
-        for inputs, targets in dataloader:
-            inputs, targets = inputs.to(device), targets.to(device)
-            logits = get_logits_from_model(model, inputs)
-            logits = logits.contiguous().view(-1, logits.size(-1))
-            targets = targets.contiguous().view(-1)[:logits.shape[0]]
-            loss = criterion(logits, targets)
-            total_loss += loss.item()
-            count += 1
-    avg_loss = total_loss / max(count, 1)
-    model.train()
-    return avg_loss
-
-# ============================= CLEANUP OLD EPOCHS =============================
-
-def cleanup_old_epochs(keep_last=KEEP_LAST_EPOCHS):
-    epochs = sorted([p for p in OUTPUT_DIR.glob("epoch*") if p.is_dir()],
-                    key=lambda x: int(x.name.replace("epoch", "")))
-    for old in epochs[:-keep_last]:
-        if old.exists():
-            shutil.rmtree(old)
-            print(f"Old epoch deleted: {old.name}")
-
-# ============================= TRAINING =============================
-
-def train():
-    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-    print("Loading model...")
-    model = None
-    if LAST_TRAINED_PATH.exists():
-        print(f"Continuing training from last JIT model: {LAST_TRAINED_PATH}")
-        model = torch.jit.load(LAST_TRAINED_PATH, map_location=device)
-    elif BASE_MODEL_PATH.exists():
-        print(f"Starting from base JIT model: {BASE_MODEL_PATH}")
-        model = torch.jit.load(BASE_MODEL_PATH, map_location=device)
-    else:
-        print(f"ERROR: JIT model not found. Checked paths: {BASE_MODEL_PATH} and {LAST_TRAINED_PATH}")
-        print("Please run the JIT export script (e.g., 'model_export.py') first.")
-        return
-
-    # Sometimes torch.jit.load with map_location doesn't move every internal buffer.
-    # Force a device move for ScriptModule, wrapped in try/except for compatibility.
-    try:
-        model.to(device)
-    except Exception:
-        # If ScriptModule.to fails, attempt moving by reloading state_dict -> module approach is not always possible.
-        pass
-
-    # As extra safety, try to move any freqs buffers inside submodules (best-effort).
-    try:
-        for name, buf in model.named_buffers():
-            if buf is not None and buf.device != device:
-                try:
-                    model.register_buffer(name, buf.to(device))
-                except Exception:
-                    # Some ScriptModule buffers may not be re-registerable; ignore non-critical failures.
-                    pass
-    except Exception:
-        pass
-
-    # Проверка весов на NaN/Inf
-    try:
-        for n, p in model.named_parameters():
-            if torch.isnan(p).any():
-                print(f"[FATAL] NaN in weights: {n}")
-                exit(10)
-            if torch.isinf(p).any():
-                print(f"[FATAL] Inf in weights: {n}")
-                exit(11)
-    except Exception:
-        # some JIT modules may not expose named_parameters() - ignore if unavailable
-        pass
-
-    model.train()
-    try:
-        model.gradient_checkpointing_enable()
-        print("✅ Gradient Checkpointing Enabled.")
-    except Exception:
-        print("⚠️ Warning: model.gradient_checkpointing_enable() not found on JIT model. Training will proceed without GC.")
-
-    # =========================================================================
-    # ФИНАЛЬНОЕ ИСПРАВЛЕНИЕ: Используем ЛОКАЛЬНУЮ ПАПКУ токенизатора
-    # =========================================================================
-    LOCAL_TOKENIZER_PATH = "./tokenizer" # Путь к папке, куда вы загрузили файлы токенизатора
-    
-    train_dataset = LazyTextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, tokenizer_name=LOCAL_TOKENIZER_PATH, split_type='train', val_ratio=VAL_SPLIT_RATIO)
-    val_dataset = LazyTextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, tokenizer_name=LOCAL_TOKENIZER_PATH, split_type='val', val_ratio=VAL_SPLIT_RATIO)
-    # =========================================================================
-    
-    # IterableDataset: must use drop_last=True and shuffle=False, num_workers=0 on CPU/GPU
-    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True, num_workers=0)
-    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True, num_workers=0)
-
-    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
-    criterion = nn.CrossEntropyLoss()
-
-    total_steps = (len(train_dataset) // BATCH_SIZE) * EPOCHS
-    print(f"\n=== BEGINNING LONG-TERM TRAINING ===")
-    print(f"Epochs: {EPOCHS} | Steps (Train): {total_steps} | Examples (Train): {len(train_dataset)}")
-    print(f"Batch Size (Effective): {BATCH_SIZE} | Precision: FP32")
-
-    global_step = 0
-    for epoch in range(1, EPOCHS + 1):
-        print(f"\n--- Epoch {epoch}/{EPOCHS} ---")
-        epoch_loss = 0.0
-
-        with tqdm(train_dataloader, desc=f"Epoch {epoch} [TRAIN]", leave=False) as pbar:
-            for inputs, targets in pbar:
-                inputs, targets = inputs.to(device), targets.to(device)
-                optimizer.zero_grad()
-                logits = get_logits_from_model(model, inputs)
-                logits = logits.contiguous().view(-1, logits.size(-1))
-                targets_view = targets.contiguous().view(-1)[:logits.shape[0]]
-                loss = criterion(logits, targets_view)
-                loss.backward()
-                try:
-                    torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
-                except Exception:
-                    pass
-                optimizer.step()
-
-                loss_val = loss.item()
-                epoch_loss += loss_val
-                global_step += 1
-
-                pbar.set_postfix({
-                    "loss": f"{loss_val:.3f}",
-                    "ppl": f"{math.exp(min(loss_val, 10)):.1f}",
-                    "step": f"{global_step}"
-                })
-
-        avg_train_loss = epoch_loss / max(1, len(train_dataset) // BATCH_SIZE)
-        print(f"  [TRAIN] Average loss: {avg_train_loss:.3f} | PPL: {math.exp(avg_train_loss):.1f}")
-
-        print("  [VALIDATION] Starting evaluation...")
-        val_loss = evaluate(model, val_dataloader, criterion, device)
-        print(f"  [VALIDATION] Average loss: {val_loss:.3f} | PPL: {math.exp(val_loss):.1f}")
-
-        epoch_dir = OUTPUT_DIR / f"epoch{epoch}"
-        epoch_dir.mkdir(exist_ok=True)
-        try:
-            torch.jit.save(model, epoch_dir / MODEL_SAVE_NAME)
-            print(f"Model saved: {epoch_dir / MODEL_SAVE_NAME}")
-        except Exception:
-            # If saving scripted model fails, fallback to state_dict
-            torch.save(model.state_dict(), epoch_dir / "state_dict.pt")
-            print(f"State dict saved: {epoch_dir / 'state_dict.pt'}")
-        cleanup_old_epochs()
-
-    final_dir = OUTPUT_DIR / "final"
-    final_dir.mkdir(exist_ok=True)
-    try:
-        torch.jit.save(model, final_dir / MODEL_SAVE_NAME)
-    except Exception:
-        torch.save(model.state_dict(), final_dir / "state_dict.pt")
-    # Try to save tokenizer if available
-    try:
-        train_dataset.tokenizer.save_pretrained(final_dir)
-    except Exception:
-        pass
-
-    if LAST_TRAINED_PATH.exists():
-        backup_path = BACKUP_DIR / f"gpt_last_trained_backup_{int(os.path.getmtime(LAST_TRAINED_PATH))}.script.pt"
-        shutil.copy(LAST_TRAINED_PATH, backup_path)
-        print(f"Backup of previous model created → {backup_path.name}")
-
-    try:
-        shutil.copy(final_dir / MODEL_SAVE_NAME, LAST_TRAINED_PATH)
-        print(f"Last trained model saved → {LAST_TRAINED_PATH}")
-    except Exception:
-        pass
-
-    print(f"\nTRAINING COMPLETED! Model ready:")
-    print(f"   • For chat: {final_dir / MODEL_SAVE_NAME}")
-    print(f"   • For further fine-tuning: {LAST_TRAINED_PATH}")
-
-if __name__ == "__main__":
-    if not RAW_PATH.exists():
-        print(f"ERROR: No file {RAW_PATH}")
-        print("Put your text into datasets/dialogues_text.txt")
-    else:
+# Copyright (c) 2025 CMS Manhattan
+# All rights reserved.
+# Author: Konstantin Vladimirovich Grabko
+# Email: grabko@cmsmanhattan.com
+# Phone: +1(516)777-0945
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, version 3 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+# Additional terms:
+# Any commercial use or distribution of this software or derivative works
+# requires explicit written permission from the copyright holder.
+
+
+"""
+Before run this script, download the GPT-2 tokenizer files into a local folder named 'tokenizer':
+mkdir -p tokenizer
+wget -O tokenizer/tokenizer.json          https://huggingface.co/gpt2/resolve/main/tokenizer.json
+wget -O tokenizer/vocab.json              https://huggingface.co/gpt2/resolve/main/vocab.json
+wget -O tokenizer/merges.txt              https://huggingface.co/gpt2/resolve/main/merges.txt
+wget -O tokenizer/tokenizer_config.json   https://huggingface.co/gpt2/resolve/main/tokenizer_config.json
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import IterableDataset, DataLoader
+from transformers import GPT2TokenizerFast
+from tqdm import tqdm
+import shutil
+import math
+from pathlib import Path
+import re
+
+# ============================= SETTINGS =============================
+TRAIN_SEQ_LEN = 256
+BATCH_SIZE = 1
+EPOCHS = 1
+LEARNING_RATE = 6e-6
+WEIGHT_DECAY = 0.01
+GRAD_CLIP = 1.0
+KEEP_LAST_EPOCHS = 3
+VAL_SPLIT_RATIO = 0.05
+VOCAB_SIZE = 50257
+
+BASE_MODEL_PATH = Path("models/gpt_modern_1b_class.script.pt")
+LAST_TRAINED_PATH = Path("models/gpt_1b_last_trained.script.pt")
+BACKUP_DIR = Path("models/backups")
+BACKUP_DIR.mkdir(exist_ok=True)
+
+RAW_PATH = Path("datasets/dialogues_text.txt")
+CLEAN_PATH = Path("datasets/dialogues_text_clean.txt")
+
+# Device selection
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# -- Dataset cleaning --
+force_clean = False
+if not CLEAN_PATH.exists():
+    print("Cleaned dataset not found. Performing initial cleaning...")
+    force_clean = True
+else:
+    try:
+        if RAW_PATH.stat().st_mtime > CLEAN_PATH.stat().st_mtime:
+            print("Detected changes in the raw dataset. Re-cleaning...")
+            force_clean = True
+        else:
+            print(f"Using existing cleaned dataset → {CLEAN_PATH}")
+    except FileNotFoundError:
+        print("File system synchronization error. Performing re-cleaning for safety...")
+        force_clean = True
+
+if force_clean:
+    if not RAW_PATH.exists():
+        raise FileNotFoundError(f"ERROR: Source file {RAW_PATH} not found. Check the path.")
+
+    print("Cleaning up the dataset from garbage (wrong separators, extra spaces)...")
+    text = RAW_PATH.read_text(encoding="utf-8")
+    text = re.sub(r' {2,}', ' ', text)
+    text = text.replace(" \n", "\n").replace("\n ", "\n")
+    CLEAN_PATH.write_text(text, encoding="utf-8")
+    print(f"Dataset successfully cleaned and saved → {CLEAN_PATH}")
+
+DATASET_PATH = CLEAN_PATH
+OUTPUT_DIR = Path("build/fine_tuning_output")
+MODEL_SAVE_NAME = "gpt_finetuned.script.pt"
+
+# ============================= DATASET (LAZY) =============================
+
+class LazyTextDataset(IterableDataset):
+    """Lazy memory-efficient dataset, splits on-the-fly into train and val."""
+    # Обратите внимание: аргумент tokenizer_name по-прежнему имеет значение по умолчанию "gpt2",
+    # но в функции train() мы теперь передаем локальный путь.
+    def __init__(self, text_file, seq_len=TRAIN_SEQ_LEN, tokenizer_name="gpt2", split_type='train', val_ratio=VAL_SPLIT_RATIO):
+        self.seq_len = seq_len
+        # Эта строка теперь загружает токенизатор из локальной папки, если передан локальный путь.
+        self.tokenizer = GPT2TokenizerFast.from_pretrained(tokenizer_name)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.text_file = text_file
+        self.split_type = split_type
+        self.val_ratio = val_ratio
+
+        print(f"Loading and tokenizing text from {text_file}")
+        with open(text_file, "r", encoding="utf-8") as f:
+            self.data = f.read()
+        self.tokens = self.tokenizer.encode(self.data)
+        
+        # Work out split indices
+        total_tokens = len(self.tokens) - 1  # because label sequence shifted
+        total_batches = total_tokens // seq_len
+        val_size = int(total_batches * self.val_ratio)
+        train_size = total_batches - val_size
+        if split_type == 'train':
+            self.start = 0
+            self.stop = train_size
+        elif split_type == 'val':
+            self.start = train_size
+            self.stop = train_size + val_size
+        else:
+            raise ValueError(f"split_type should be 'train' or 'val', got {split_type}")
+        self.total_sequences = self.stop - self.start
+        print(f"Lazy dataset: {self.total_sequences:,} sequences for {split_type} split (from {total_batches:,} total)")
+
+    def __iter__(self):
+        for i in range(self.start * self.seq_len, self.stop * self.seq_len, self.seq_len):
+            # Make sure last batch fits
+            if i + self.seq_len + 1 > len(self.tokens):
+                break
+            input_seq = torch.tensor(self.tokens[i : i + self.seq_len], dtype=torch.long)
+            label_seq = torch.tensor(self.tokens[i + 1 : i + self.seq_len + 1], dtype=torch.long)
+            yield input_seq, label_seq
+
+    def __len__(self):
+        return self.total_sequences
+
+# ============================= GET LOGITS UTIL =============================
+
+def get_logits_from_model(model, inputs):
+    """
+    Robust wrapper to call either a scripted JIT model or nn.Module.
+    Handles models that either return (logits, kv) or just logits.
+    """
+    # Ensure inputs on same device as model parameters/buffers
+    inputs = inputs.to(device)
+    try:
+        out = model(inputs)
+    except RuntimeError as e:
+        # Some JIT modules expect plain tensor on CPU device for tracing path.
+        # Re-raise if unrelated
+        raise
+
+    # Model may return logits or (logits, kv)
+    if isinstance(out, tuple) or (isinstance(out, list) and len(out) >= 1):
+        logits = out[0]
+    else:
+        logits = out
+    return logits
+
+# ============================= EVALUATION (VALIDATION) =============================
+
+def evaluate(model, dataloader, criterion, device):
+    model.eval()
+    total_loss = 0.0
+    count = 0
+    with torch.no_grad():
+        for inputs, targets in dataloader:
+            inputs, targets = inputs.to(device), targets.to(device)
+            logits = get_logits_from_model(model, inputs)
+            logits = logits.contiguous().view(-1, logits.size(-1))
+            targets = targets.contiguous().view(-1)[:logits.shape[0]]
+            loss = criterion(logits, targets)
+            total_loss += loss.item()
+            count += 1
+    avg_loss = total_loss / max(count, 1)
+    model.train()
+    return avg_loss
+
+# ============================= CLEANUP OLD EPOCHS =============================
+
+def cleanup_old_epochs(keep_last=KEEP_LAST_EPOCHS):
+    epochs = sorted([p for p in OUTPUT_DIR.glob("epoch*") if p.is_dir()],
+                    key=lambda x: int(x.name.replace("epoch", "")))
+    for old in epochs[:-keep_last]:
+        if old.exists():
+            shutil.rmtree(old)
+            print(f"Old epoch deleted: {old.name}")
+
+# ============================= TRAINING =============================
+
+def train():
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    print("Loading model...")
+    model = None
+    if LAST_TRAINED_PATH.exists():
+        print(f"Continuing training from last JIT model: {LAST_TRAINED_PATH}")
+        model = torch.jit.load(LAST_TRAINED_PATH, map_location=device)
+    elif BASE_MODEL_PATH.exists():
+        print(f"Starting from base JIT model: {BASE_MODEL_PATH}")
+        model = torch.jit.load(BASE_MODEL_PATH, map_location=device)
+    else:
+        print(f"ERROR: JIT model not found. Checked paths: {BASE_MODEL_PATH} and {LAST_TRAINED_PATH}")
+        print("Please run the JIT export script (e.g., 'model_export.py') first.")
+        return
+
+    # Sometimes torch.jit.load with map_location doesn't move every internal buffer.
+    # Force a device move for ScriptModule, wrapped in try/except for compatibility.
+    try:
+        model.to(device)
+    except Exception:
+        # If ScriptModule.to fails, attempt moving by reloading state_dict -> module approach is not always possible.
+        pass
+
+    # As extra safety, try to move any freqs buffers inside submodules (best-effort).
+    try:
+        for name, buf in model.named_buffers():
+            if buf is not None and buf.device != device:
+                try:
+                    model.register_buffer(name, buf.to(device))
+                except Exception:
+                    # Some ScriptModule buffers may not be re-registerable; ignore non-critical failures.
+                    pass
+    except Exception:
+        pass
+
+    # Проверка весов на NaN/Inf
+    try:
+        for n, p in model.named_parameters():
+            if torch.isnan(p).any():
+                print(f"[FATAL] NaN in weights: {n}")
+                exit(10)
+            if torch.isinf(p).any():
+                print(f"[FATAL] Inf in weights: {n}")
+                exit(11)
+    except Exception:
+        # some JIT modules may not expose named_parameters() - ignore if unavailable
+        pass
+
+    model.train()
+    try:
+        model.gradient_checkpointing_enable()
+        print("✅ Gradient Checkpointing Enabled.")
+    except Exception:
+        print("⚠️ Warning: model.gradient_checkpointing_enable() not found on JIT model. Training will proceed without GC.")
+
+    # =========================================================================
+    # ФИНАЛЬНОЕ ИСПРАВЛЕНИЕ: Используем ЛОКАЛЬНУЮ ПАПКУ токенизатора
+    # =========================================================================
+    LOCAL_TOKENIZER_PATH = "./tokenizer" # Путь к папке, куда вы загрузили файлы токенизатора
+    
+    train_dataset = LazyTextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, tokenizer_name=LOCAL_TOKENIZER_PATH, split_type='train', val_ratio=VAL_SPLIT_RATIO)
+    val_dataset = LazyTextDataset(DATASET_PATH, seq_len=TRAIN_SEQ_LEN, tokenizer_name=LOCAL_TOKENIZER_PATH, split_type='val', val_ratio=VAL_SPLIT_RATIO)
+    # =========================================================================
+    
+    # IterableDataset: must use drop_last=True and shuffle=False, num_workers=0 on CPU/GPU
+    train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True, num_workers=0)
+    val_dataloader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, drop_last=True, num_workers=0)
+
+    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
+    criterion = nn.CrossEntropyLoss()
+
+    total_steps = (len(train_dataset) // BATCH_SIZE) * EPOCHS
+    print(f"\n=== BEGINNING LONG-TERM TRAINING ===")
+    print(f"Epochs: {EPOCHS} | Steps (Train): {total_steps} | Examples (Train): {len(train_dataset)}")
+    print(f"Batch Size (Effective): {BATCH_SIZE} | Precision: FP32")
+
+    global_step = 0
+    for epoch in range(1, EPOCHS + 1):
+        print(f"\n--- Epoch {epoch}/{EPOCHS} ---")
+        epoch_loss = 0.0
+
+        with tqdm(train_dataloader, desc=f"Epoch {epoch} [TRAIN]", leave=False) as pbar:
+            for inputs, targets in pbar:
+                inputs, targets = inputs.to(device), targets.to(device)
+                optimizer.zero_grad()
+                logits = get_logits_from_model(model, inputs)
+                logits = logits.contiguous().view(-1, logits.size(-1))
+                targets_view = targets.contiguous().view(-1)[:logits.shape[0]]
+                loss = criterion(logits, targets_view)
+                loss.backward()
+                try:
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+                except Exception:
+                    pass
+                optimizer.step()
+
+                loss_val = loss.item()
+                epoch_loss += loss_val
+                global_step += 1
+
+                pbar.set_postfix({
+                    "loss": f"{loss_val:.3f}",
+                    "ppl": f"{math.exp(min(loss_val, 10)):.1f}",
+                    "step": f"{global_step}"
+                })
+
+        avg_train_loss = epoch_loss / max(1, len(train_dataset) // BATCH_SIZE)
+        print(f"  [TRAIN] Average loss: {avg_train_loss:.3f} | PPL: {math.exp(avg_train_loss):.1f}")
+
+        print("  [VALIDATION] Starting evaluation...")
+        val_loss = evaluate(model, val_dataloader, criterion, device)
+        print(f"  [VALIDATION] Average loss: {val_loss:.3f} | PPL: {math.exp(val_loss):.1f}")
+
+        epoch_dir = OUTPUT_DIR / f"epoch{epoch}"
+        epoch_dir.mkdir(exist_ok=True)
+        try:
+            torch.jit.save(model, epoch_dir / MODEL_SAVE_NAME)
+            print(f"Model saved: {epoch_dir / MODEL_SAVE_NAME}")
+        except Exception:
+            # If saving scripted model fails, fallback to state_dict
+            torch.save(model.state_dict(), epoch_dir / "state_dict.pt")
+            print(f"State dict saved: {epoch_dir / 'state_dict.pt'}")
+        cleanup_old_epochs()
+
+    final_dir = OUTPUT_DIR / "final"
+    final_dir.mkdir(exist_ok=True)
+    try:
+        torch.jit.save(model, final_dir / MODEL_SAVE_NAME)
+    except Exception:
+        torch.save(model.state_dict(), final_dir / "state_dict.pt")
+    # Try to save tokenizer if available
+    try:
+        train_dataset.tokenizer.save_pretrained(final_dir)
+    except Exception:
+        pass
+
+    if LAST_TRAINED_PATH.exists():
+        backup_path = BACKUP_DIR / f"gpt_last_trained_backup_{int(os.path.getmtime(LAST_TRAINED_PATH))}.script.pt"
+        shutil.copy(LAST_TRAINED_PATH, backup_path)
+        print(f"Backup of previous model created → {backup_path.name}")
+
+    try:
+        shutil.copy(final_dir / MODEL_SAVE_NAME, LAST_TRAINED_PATH)
+        print(f"Last trained model saved → {LAST_TRAINED_PATH}")
+    except Exception:
+        pass
+
+    print(f"\nTRAINING COMPLETED! Model ready:")
+    print(f"   • For chat: {final_dir / MODEL_SAVE_NAME}")
+    print(f"   • For further fine-tuning: {LAST_TRAINED_PATH}")
+
+if __name__ == "__main__":
+    if not RAW_PATH.exists():
+        print(f"ERROR: No file {RAW_PATH}")
+        print("Put your text into datasets/dialogues_text.txt")
+    else:
         train()