#%%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()