train_jirack_accelerate.py

#%%writefile  train_jirack_accelerate_v3.py

# =============================================================================
# COPYRIGHT © 2025-2026  Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
# CMS Manhattan JiRack Technology — PATENT PENDING
#
# This code is proprietary. 
# Personal and non-commercial research use is allowed.
# Any commercial use, derivative works for profit, or distribution 
# requires a paid license and 5% royalty.
#
# Unauthorized commercial use is strictly prohibited.
# Contact: grabko@cmsmanhattan.com
# =============================================================================
import os
# Включаем оптимизацию памяти для ROCm/HIP ДО импорта torch!
os.environ["PYTORCH_HIP_ALLOC_CONF"] = "expandable_segments:True"

import re
import glob
import math
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from accelerate import Accelerator
from tqdm import tqdm
from transformers import Adafactor, get_cosine_schedule_with_warmup

# --- GLOBAL CONSTANTS ---
USE_COSINE_SCHEDULER = False  # Set to False if you want to disable the scheduler (warmup + cosine)

# --- 1. Lightweight Dataset ---
class SingleShardDataset(Dataset):
    def __init__(self, shard_path):
        self.data = torch.load(shard_path, map_location="cpu", weights_only=True)
    def __len__(self):
        return self.data.shape[0]
    def __getitem__(self, idx):
        return self.data[idx].long()

# --- Helpers for Natural Sorting ---
def natural_sort_key(s):
    return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', s)]

def extract_shard_number(filename):
    match = re.search(r'model_weights_shard_(\d+)\.pt', filename)
    return int(match.group(1)) if match else 0

# --- 2. Main Training Function ---
def train():
    grad_accumulation_steps = 24
    batch_size = 8

    # Initialize Accelerator
    accelerator = Accelerator(
        mixed_precision="bf16",
        gradient_accumulation_steps=grad_accumulation_steps
    )

    # Определяем абсолютные пути относительно расположения самого скрипта
    script_dir = os.path.dirname(os.path.abspath(__file__))
    pt_chunks_mask = os.path.join(script_dir, "pretraindata/jirack_pretrain_chunk_*.pt")
    checkpoint_dir = os.path.join(script_dir, "checkpoints")
    
    pt_files = sorted(glob.glob(pt_chunks_mask), key=natural_sort_key)

    if not pt_files:
        raise FileNotFoundError(f"No chunk files found matching the mask: {pt_chunks_mask}")

    if accelerator.is_local_main_process:
        print(f"Shards found: {len(pt_files)}")
        print("Initializing JiRack 3.3B...")

    # Импортируем строго ваши классы из локального файла
    from JiRackNative_3b import TernaryTransformer3B, TernaryConfig

    config = TernaryConfig()
    model = TernaryTransformer3B(config)

    # === DYNAMIC CHECKPOINT DISCOVERY (ABSOLUTE PATHS) ===
    checkpoint_load_path = os.path.join(script_dir, "model_weights.pt")
    start_shard_idx = 0

    # Сканируем папку checkpoints по абсолютному пути
    shard_checkpoints = glob.glob(os.path.join(checkpoint_dir, "model_weights_shard_*.pt"))
    
    if shard_checkpoints:
        shard_checkpoints = sorted(shard_checkpoints, key=natural_sort_key)
        latest_shard_checkpoint = shard_checkpoints[-1]
        completed_shards = extract_shard_number(latest_shard_checkpoint)
        
        checkpoint_load_path = latest_shard_checkpoint
        start_shard_idx = completed_shards  # Если закончили шард 2, индекс следующего равен 2 (3-й шард)

    # Загружаем найденные веса
    if os.path.exists(checkpoint_load_path):
        if accelerator.is_local_main_process:
            print(f"-> Loading saved weights from: {checkpoint_load_path}")
        state_dict = torch.load(checkpoint_load_path, map_location="cpu", weights_only=True)
        model.load_state_dict(state_dict)
        if accelerator.is_local_main_process and start_shard_idx > 0:
            print(f"-> Resuming training from Shard {start_shard_idx + 1} (Skipping first {start_shard_idx} shards)")
    else:
        if accelerator.is_local_main_process:
            print(f"-> Checkpoint not found at {checkpoint_load_path}, training will start from scratch.")

    model.gradient_checkpointing_enable()

    if accelerator.is_local_main_process:
        print("-> Gradient Checkpointing enabled.")

    criterion = nn.CrossEntropyLoss()
    model = accelerator.prepare(model)

    # === ADAFACTOR CONFIGURATION FOR GPU ===
    optimizer = Adafactor(
        model.parameters(),
        lr=2e-4,
        weight_decay=0.01,
        relative_step=False,
        scale_parameter=False,
        warmup_init=False
    )

    # === CALCULATING AND INITIALIZING SCHEDULER WITH WARMUP ===
    scheduler = None
    if USE_COSINE_SCHEDULER:
        steps_per_shard = math.ceil(2000 / (batch_size * grad_accumulation_steps))
        total_steps = len(pt_files) * steps_per_shard
        num_warmup_steps = int(0.05 * total_steps)

        scheduler = get_cosine_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=num_warmup_steps,
            num_training_steps=total_steps
        )

        if accelerator.is_local_main_process:
            print(f"-> Scheduler ACTIVATED. Total steps: {total_steps}")

    if scheduler is not None:
        optimizer, scheduler = accelerator.prepare(optimizer, scheduler)
    else:
        optimizer = accelerator.prepare(optimizer)

    model.train()
    if accelerator.is_local_main_process:
        print("Starting training...")
        os.makedirs(checkpoint_dir, exist_ok=True)

    # === ОСНОВНОЙ ЦИКЛ ПО ШАРДАМ ===
    for shard_counter, shard_path in enumerate(pt_files):
        # Пропускаем уже обработанные шарды
        if shard_counter < start_shard_idx:
            continue

        shard_name = os.path.basename(shard_path)
        if accelerator.is_local_main_process:
            print(f"\n[Shard {shard_counter + 1}/{len(pt_files)}] {shard_name}")

        shard_dataset = SingleShardDataset(shard_path)

        train_loader = DataLoader(
            shard_dataset,
            batch_size=batch_size,
            shuffle=True
        )

        train_loader = accelerator.prepare(train_loader)

        progress_bar = tqdm(
            train_loader,
            desc=f"Training {shard_name}",
            disable=not accelerator.is_local_main_process
        )

        epoch_loss = 0.0
        for step, batch in enumerate(progress_bar):
            input_ids = batch
            inputs = input_ids[:, :-1]
            targets = input_ids[:, 1:]

            with accelerator.accumulate(model):
                logits, _ = model(inputs)
                loss = criterion(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))

                accelerator.backward(loss)
                optimizer.step()

                if scheduler is not None and accelerator.sync_gradients:
                    if not getattr(accelerator, "optimizer_step_was_skipped", False):
                        scheduler.step()

                optimizer.zero_grad()

                epoch_loss += loss.item()
                avg_loss = epoch_loss / (step + 1)
                ppl = math.exp(avg_loss) if avg_loss < 20 else float('inf')

                if accelerator.is_local_main_process:
                    current_lr = scheduler.get_last_lr()[0] if scheduler is not None else optimizer.param_groups[0]['lr']
                    progress_bar.set_postfix({
                        "loss": f"{loss.item():.4f}",
                        "avg_loss": f"{avg_loss:.4f}",
                        "ppl": f"{ppl:.1f}",
                        "lr": f"{current_lr:.2e}"
                    })

        # Сохраняем веса с абсолютным путем
        if accelerator.is_local_main_process:
            shard_checkpoint_path = os.path.join(checkpoint_dir, f"model_weights_shard_{shard_counter + 1}.pt")
            unwrapped_model = accelerator.unwrap_model(model)
            torch.save(unwrapped_model.state_dict(), shard_checkpoint_path)
            print(f"✅   Saved after shard {shard_counter + 1}: {shard_checkpoint_path}")

    # Финальное сохранение
    accelerator.wait_for_everyone()
    if accelerator.is_local_main_process:
        final_path = os.path.join(checkpoint_dir, "model_final_weights.pt")
        unwrapped_model = accelerator.unwrap_model(model)
        torch.save(unwrapped_model.state_dict(), final_path)
        print(f"Final save completed: {final_path}")

if __name__ == "__main__":
    train()