Add files using upload-large-folder tool

.gitattributes

@@ -35,3 +35,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 train_hnet_with_docstring_18_04/wandb/run-20260417_085757-sa79g3yl/run-sa79g3yl.wandb filter=lfs diff=lfs merge=lfs -text
 wandb/run-20260418_121916-2mk39j3k/run-2mk39j3k.wandb filter=lfs diff=lfs merge=lfs -text
+pythia1b_v5_04_21/wandb/run-20260421_202839-8ing6xdi/run-8ing6xdi.wandb filter=lfs diff=lfs merge=lfs -text

pythia1b_v5_04_21/checkpoints/checkpoint_latest.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b58e3452e25fa6a1b168abc2eeb93bfd178bd0d36dedbb37170f173dd7a163d
+size 6070947158

pythia1b_v5_04_21/checkpoints/checkpoint_step_10591.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b82966ae12c1e00b13902b6d602e024a5131f8f51b2de7fcfdd95b689050729
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_12000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a35b05f715bb652fa11b55de09cf91de2bb1cfd86ebd9c6dedb708608a0d412c
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_15000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44f0611f7af451e017f865eeaa0281b81ed3a02071e89003d9b5b6e95a1b19d5
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_18000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6cfdf9707733ebf03847f06688369930cc98a11842515a343355432fd5f1b521
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_21000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d4ac79d33bc5ef5a51922f227e04f969c5b69356dc854730897d98bb10d9e731
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_21182.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:505dcf26777c122f00bd156aed99f953998a4ec32972fbd3d091ebbe26ee0be4
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_24000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:75521e0fbe7d16afd7a96599a0cbb6ed91ab8f827cf4b8ebe194e89959ec16a1
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_27000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85a288ed13b41ef7c4d6b185eae644005f8e4f0dbcb930a8ad620668ffcd9c37
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_3000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c372dfd8cc37b45a6747878c0f1ca532c2775264c3814cc5e77edc85b6cf1553
+size 6070949586

pythia1b_v5_04_21/checkpoints/checkpoint_step_30000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5514bd8b466d5dfbe060bd1030688b3e8631a3a7cc74341f797fcb81bb07202
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_31773.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23a2f1b814172d824cc94fe4eb671efb85eddc0d701dc770bbd2510a862dc9f2
+size 6070950374

pythia1b_v5_04_21/checkpoints/checkpoint_step_6000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99b4d33f41190374d0a57e12b5715a87ca53f4693512fe0a8d518ff852ccc7e7
+size 6070949586

pythia1b_v5_04_21/checkpoints/checkpoint_step_9000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f35f51cf37687952cc2753d63e1d0ccfd173b4f43a345440f98613c07c39ac5
+size 6070949586

pythia1b_v5_04_21/model_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:777b8524d357a844e33c810453f9bca93a43ee31f69a72619ba2b24f61487d48
+size 2023640386

pythia1b_v5_04_21/model_final.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:738b4bf2c5af646ac8c184cc78e3a2359d1a6689362cbca0b53ab0cd7d2fea53
+size 2023640586

pythia1b_v5_04_21/wandb/run-20260421_202839-8ing6xdi/files/code/code_completion_exp/train_pythia/train.py

@@ -0,0 +1,598 @@
+"""
+Training Pipeline для Pythia (decoder-only transformer) на задаче Code Completion.
+
+Конфигурация через Hydra + OmegaConf, логирование в Trackio.
+Поддержка DDP через Accelerate для multi-GPU тренировки.
+
+Использование:
+    # Базовый запуск (single GPU)
+    python train.py
+
+    # Multi-GPU с Accelerate
+    accelerate launch train.py
+
+    # Multi-GPU с указанием количества GPU
+    accelerate launch --num_processes=4 train.py
+
+    # Переопределение параметров через CLI
+    python train.py training.lr=1e-4 training.epochs=5
+
+    # Выбор другого конфига модели
+    python train.py model=pythia_160m
+
+    # Multirun (sweep)
+    python train.py --multirun training.lr=1e-4,3e-4,1e-3
+
+    # Без логирования
+    python train.py tracking.enabled=false
+"""
+
+import os
+import math
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from datasets import load_from_disk
+
+import hydra
+from hydra.core.hydra_config import HydraConfig
+from omegaconf import DictConfig, OmegaConf
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoConfig,
+    PreTrainedTokenizerBase,
+)
+from accelerate import Accelerator
+from accelerate.utils import set_seed as accelerate_set_seed
+
+# Ensure repo root is on sys.path (needed when running from subdirectory)
+import sys
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+# Shared training library
+from training_lib.utils import AverageMeter, log_message
+from training_lib.checkpointing import save_checkpoint, load_checkpoint
+from training_lib.schedulers import get_lr_scheduler
+from training_lib.tracking import init_tracking, log_metrics, finish_tracking
+from training_lib.validation import run_validation
+
+
+# ============================================================================
+# ДАННЫЕ
+# ============================================================================
+
+
+class CodeCompletionCollator:
+    """Collate function для батчирования примеров code completion."""
+
+    def __init__(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        max_context_len: int = 1024,
+        max_target_len: int = 256,
+    ):
+        self.tokenizer = tokenizer
+        self.max_context_len = max_context_len
+        self.max_target_len = max_target_len
+        self.pad_token_id = tokenizer.pad_token_id
+
+    def __call__(self, batch: list[dict]) -> dict:
+        contexts = [item["context"] for item in batch]
+        targets = [item["target"] for item in batch]
+
+        encoded_contexts = self.tokenizer(
+            contexts,
+            add_special_tokens=True,
+            truncation=True,
+            max_length=self.max_context_len,
+            return_tensors=None,
+        )
+        encoded_targets = self.tokenizer(
+            targets,
+            add_special_tokens=False,
+            truncation=True,
+            max_length=self.max_target_len,
+            return_tensors=None,
+        )
+
+        input_ids_list = []
+        context_lengths = []
+
+        for ctx_ids, tgt_ids in zip(
+            encoded_contexts["input_ids"], encoded_targets["input_ids"]
+        ):
+            tgt_ids = tgt_ids + [self.tokenizer.eos_token_id]
+            context_lengths.append(len(ctx_ids))
+            input_ids_list.append(ctx_ids + tgt_ids)
+
+        max_len = max(len(ids) for ids in input_ids_list)
+
+        padded_input_ids = []
+        attention_mask = []
+
+        for ids in input_ids_list:
+            padding_len = max_len - len(ids)
+            padded_input_ids.append(ids + [self.pad_token_id] * padding_len)
+            attention_mask.append([1] * len(ids) + [0] * padding_len)
+
+        return {
+            "input_ids": torch.tensor(padded_input_ids, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+            "context_lengths": torch.tensor(context_lengths, dtype=torch.long),
+        }
+
+
+def create_dataloaders(
+    cfg: DictConfig, tokenizer: PreTrainedTokenizerBase
+) -> dict[str, DataLoader]:
+    """Создание DataLoader'ов для train и validation."""
+    dataset_dict = load_from_disk(cfg.data.path)
+
+    collator = CodeCompletionCollator(
+        tokenizer=tokenizer,
+        max_context_len=cfg.data.max_context_len,
+        max_target_len=cfg.data.max_target_len,
+    )
+
+    dataloaders = {}
+
+    if "train" in dataset_dict:
+        dataloaders["train"] = DataLoader(
+            dataset_dict["train"],
+            batch_size=cfg.training.batch_size,
+            shuffle=True,
+            collate_fn=collator,
+            num_workers=cfg.data.num_workers,
+            pin_memory=cfg.data.pin_memory,
+        )
+
+    if "validation" in dataset_dict:
+        eval_batch_size = cfg.training.get("eval_batch_size", cfg.training.batch_size)
+        dataloaders["validation"] = DataLoader(
+            dataset_dict["validation"],
+            batch_size=eval_batch_size,
+            shuffle=False,
+            collate_fn=collator,
+            num_workers=cfg.data.num_workers,
+            pin_memory=cfg.data.pin_memory,
+        )
+
+    return dataloaders
+
+
+
+
+# ============================================================================
+# LOSS ФУНКЦИИ
+# ============================================================================
+
+
+def compute_loss(
+    logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    context_lengths: torch.Tensor,
+    attention_mask: torch.Tensor,
+) -> dict:
+    """Вычисление loss для авторегрессионной модели."""
+    batch_size, seq_len, vocab_size = logits.shape
+
+    shift_logits = logits[:, :-1, :].contiguous()
+    shift_labels = input_ids[:, 1:].contiguous()
+    shift_mask = attention_mask[:, 1:].contiguous()
+
+    target_mask = torch.zeros_like(shift_labels, dtype=torch.bool)
+    for i in range(batch_size):
+        ctx_len = context_lengths[i].item()
+        target_mask[i, ctx_len - 1 :] = True
+
+    final_mask = target_mask & shift_mask.bool()
+
+    if final_mask.sum() > 0:
+        loss = F.cross_entropy(
+            shift_logits[final_mask], shift_labels[final_mask], reduction="mean"
+        )
+    else:
+        loss = torch.tensor(0.0, device=logits.device)
+
+    return {"loss": loss}
+
+
+def _pythia_forward_loss(
+    model: nn.Module,
+    batch: dict,
+    cfg: DictConfig,
+    accelerator: Accelerator,
+) -> dict:
+    """Forward + loss for a plain HF causal LM (attention_mask= kwarg, .logits)."""
+    input_ids = batch["input_ids"]
+    attention_mask = batch["attention_mask"]
+    context_lengths = batch["context_lengths"]
+    output = model(input_ids, attention_mask=attention_mask)
+    return compute_loss(output.logits, input_ids, context_lengths, attention_mask)
+
+
+# ============================================================================
+# PARAMETER GROUPING
+# ============================================================================
+
+
+def group_params(model: nn.Module, weight_decay: float) -> list[dict]:
+    """Группировка параметров для optimizer."""
+    decay_params = []
+    no_decay_params = []
+    
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+            
+        if "bias" in name or "LayerNorm" in name or "layernorm" in name:
+            no_decay_params.append(param)
+        else:
+            decay_params.append(param)
+    
+    return [
+        {"params": decay_params, "weight_decay": weight_decay},
+        {"params": no_decay_params, "weight_decay": 0.0},
+    ]
+
+
+
+
+# ============================================================================
+# TRAINING LOOP
+# ============================================================================
+
+
+def train_epoch(
+    model: nn.Module,
+    dataloader: DataLoader,
+    optimizer: torch.optim.Optimizer,
+    scheduler,
+    cfg: DictConfig,
+    epoch: int,
+    global_step: int,
+    accelerator: Accelerator,
+    val_dataloader: DataLoader | None = None,
+    best_val_loss: float = float("inf"),
+) -> tuple[int, float]:
+    """Один epoch тренировки. Возвращает (global_step, best_val_loss)."""
+    model.train()
+
+    loss_meter = AverageMeter()
+
+    optimizer.zero_grad()
+    accumulated_loss = 0.0
+    accumulated_steps = 0
+
+    epoch_start_time = time.time()
+    step_start_time = time.time()
+
+    for batch_idx, batch in enumerate(dataloader):
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        context_lengths = batch["context_lengths"]
+
+        with accelerator.autocast():
+            output = model(input_ids, attention_mask=attention_mask)
+            logits = output.logits
+            loss_dict = compute_loss(
+                logits, input_ids, context_lengths, attention_mask
+            )
+
+        loss = loss_dict["loss"] / cfg.training.gradient_accumulation_steps
+        accelerator.backward(loss)
+
+        accumulated_loss += loss_dict["loss"].item()
+        accumulated_steps += 1
+
+        if accumulated_steps == cfg.training.gradient_accumulation_steps:
+            if cfg.training.max_grad_norm > 0:
+                accelerator.clip_grad_norm_(
+                    model.parameters(), cfg.training.max_grad_norm
+                )
+
+            optimizer.step()
+            scheduler.step()
+            optimizer.zero_grad()
+
+            avg_loss = accumulated_loss / cfg.training.gradient_accumulation_steps
+            loss_meter.update(avg_loss)
+
+            global_step += 1
+
+            if global_step % cfg.logging.log_interval == 0:
+                step_time = time.time() - step_start_time
+                current_lr = scheduler.get_last_lr()[0]
+
+                metrics = {
+                    "train/loss": loss_meter.val,
+                    "train/loss_avg": loss_meter.avg,
+                    "train/lr": current_lr,
+                    "train/epoch": epoch,
+                    "train/step_time": step_time / cfg.logging.log_interval,
+                }
+
+                log_metrics(metrics, step=global_step)
+
+                log_message(
+                    f"Epoch {epoch} | Step {global_step} | "
+                    f"Loss: {loss_meter.avg:.4f} | "
+                    f"LR: {current_lr:.2e}",
+                    cfg,
+                    accelerator,
+                )
+
+                step_start_time = time.time()
+
+            if (
+                cfg.logging.save_interval > 0
+                and global_step % cfg.logging.save_interval == 0
+            ):
+                save_checkpoint(
+                    model, optimizer, scheduler, global_step, epoch, cfg, accelerator
+                )
+
+            eval_interval = cfg.logging.get("eval_interval", 0)
+            if (
+                eval_interval > 0
+                and val_dataloader is not None
+                and global_step % eval_interval == 0
+            ):
+                val_metrics = run_validation(
+                    model=model,
+                    dataloader=val_dataloader,
+                    cfg=cfg,
+                    global_step=global_step,
+                    accelerator=accelerator,
+                    forward_loss_fn=_pythia_forward_loss,
+                )
+
+                if val_metrics["val/loss"] < best_val_loss:
+                    best_val_loss = val_metrics["val/loss"]
+                    if accelerator.is_main_process:
+                        best_model_path = Path(cfg.paths.output_dir) / "model_best.pt"
+                        unwrapped_model = accelerator.unwrap_model(model)
+                        torch.save(unwrapped_model.state_dict(), best_model_path)
+                        log_message(
+                            f"New best model saved! Val loss: {best_val_loss:.4f}", 
+                            cfg, 
+                            accelerator
+                        )
+
+                        log_metrics(
+                            {
+                                "best/val_loss": best_val_loss,
+                                "best/val_perplexity": val_metrics["val/perplexity"],
+                                "best/step": global_step,
+                            },
+                            step=global_step,
+                        )
+
+                model.train()
+
+            accumulated_loss = 0.0
+            accumulated_steps = 0
+
+    epoch_time = time.time() - epoch_start_time
+    
+    log_message(
+        f"Epoch {epoch} completed in {epoch_time:.2f}s | "
+        f"Loss: {loss_meter.avg:.4f}",
+        cfg,
+        accelerator,
+    )
+
+    log_metrics({
+        "epoch/loss": loss_meter.avg,
+        "epoch/time": epoch_time,
+    })
+
+    return global_step, best_val_loss
+
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+
+@hydra.main(version_base=None, config_path="configs", config_name="config")
+def main(cfg: DictConfig):
+    """Главная функция тренировки с поддержкой DDP через Accelerate."""
+
+    # === Performance: Enable TF32 for faster matmuls on Ampere+ GPUs ===
+    torch.set_float32_matmul_precision('high')
+
+    # === Accelerator Setup ===
+    mixed_precision = "bf16" if cfg.training.use_amp else "no"
+    
+    accelerator = Accelerator(
+        mixed_precision=mixed_precision,
+        gradient_accumulation_steps=cfg.training.gradient_accumulation_steps,
+    )
+
+    # === Setup ===
+    accelerate_set_seed(cfg.seed)
+    
+    if cfg.paths.output_dir is None:
+        cfg.paths.output_dir = HydraConfig.get().runtime.output_dir
+
+    OmegaConf.resolve(cfg)
+
+    log_message(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'not set')}", cfg, accelerator)
+    log_message(f"Number of processes: {accelerator.num_processes}", cfg, accelerator)
+    log_message(f"Process index: {accelerator.process_index}", cfg, accelerator)
+    log_message(f"Mixed precision: {mixed_precision}", cfg, accelerator)
+    
+    log_message("=" * 60, cfg, accelerator)
+    log_message("Pythia Training Pipeline (Hydra + Trackio + Accelerate)", cfg, accelerator)
+    log_message("=" * 60, cfg, accelerator)
+    log_message(f"Config:\n{OmegaConf.to_yaml(cfg)}", cfg, accelerator)
+
+    # === Trackio Init ===
+    init_tracking(cfg, accelerator)
+
+    # === Tokenizer ===
+    log_message("Initializing tokenizer...", cfg, accelerator)
+    tokenizer = AutoTokenizer.from_pretrained(cfg.model.name)
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    # === Model ===
+    log_message("Loading model...", cfg, accelerator)
+    
+    # Flash Attention 2
+    torch_dtype = torch.bfloat16 if cfg.training.use_amp else torch.float32
+    
+    if cfg.model.checkpoint_path:
+        model = AutoModelForCausalLM.from_pretrained(
+            cfg.model.name,
+            attn_implementation="flash_attention_2",
+            torch_dtype=torch_dtype,
+        )
+        checkpoint = torch.load(cfg.model.checkpoint_path, map_location="cpu")
+        model.load_state_dict(checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint)
+        log_message(f"Loaded checkpoint: {cfg.model.checkpoint_path}", cfg, accelerator)
+    elif cfg.model.from_scratch:
+        config = AutoConfig.from_pretrained(cfg.model.name)
+        config._attn_implementation = "flash_attention_2"
+        model = AutoModelForCausalLM.from_config(config, torch_dtype=torch_dtype)
+        log_message(f"Initialized from scratch: {cfg.model.name}", cfg, accelerator)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            cfg.model.name,
+            attn_implementation="flash_attention_2",
+            torch_dtype=torch_dtype,
+        )
+        log_message(f"Loaded pretrained: {cfg.model.name}", cfg, accelerator)
+
+    model.train()
+
+    # Log model info
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    log_message(f"Total params: {total_params:,}", cfg, accelerator)
+    log_message(f"Trainable params: {trainable_params:,}", cfg, accelerator)
+
+    # === Data ===
+    log_message("Creating dataloaders...", cfg, accelerator)
+    dataloaders = create_dataloaders(cfg, tokenizer)
+
+    train_dataloader = dataloaders["train"]
+    val_dataloader = dataloaders.get("validation", None)
+
+    log_message(f"Train dataset size: {len(train_dataloader.dataset)}", cfg, accelerator)
+    log_message(f"Train batches per epoch (before DDP split): {len(train_dataloader)}", cfg, accelerator)
+
+    if val_dataloader:
+        log_message(f"Validation dataset size: {len(val_dataloader.dataset)}", cfg, accelerator)
+        log_message(f"Validation batches: {len(val_dataloader)}", cfg, accelerator)
+    else:
+        log_message("No validation dataset found", cfg, accelerator)
+
+    # === Optimizer ===
+    log_message("Creating optimizer...", cfg, accelerator)
+    param_groups = group_params(model, cfg.training.weight_decay)
+
+    optimizer = torch.optim.AdamW(
+        param_groups,
+        lr=cfg.training.lr,
+        betas=tuple(cfg.training.betas),
+        eps=cfg.training.eps,
+    )
+
+    # === Scheduler ===
+    steps_per_epoch = math.ceil(
+        len(train_dataloader) / accelerator.num_processes
+    )
+    total_steps = (
+        cfg.training.epochs
+        * steps_per_epoch
+        // cfg.training.gradient_accumulation_steps
+    )
+    scheduler = get_lr_scheduler(optimizer, cfg, total_steps)
+
+    log_message(
+        f"Total steps: {total_steps}, Steps per epoch: {steps_per_epoch}", 
+        cfg, 
+        accelerator
+    )
+
+    # === Accelerate Prepare ===
+    log_message("Preparing model, optimizer, and dataloaders with Accelerate...", cfg, accelerator)
+    
+    if val_dataloader is not None:
+        model, optimizer, train_dataloader, val_dataloader, scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, val_dataloader, scheduler
+        )
+    else:
+        model, optimizer, train_dataloader, scheduler = accelerator.prepare(
+            model, optimizer, train_dataloader, scheduler
+        )
+
+    log_message(f"Train batches per epoch (after DDP split): {len(train_dataloader)}", cfg, accelerator)
+
+    # === Resume ===
+    global_step = 0
+    start_epoch = 1
+
+    if cfg.training.resume and cfg.training.resume_checkpoint:
+        global_step, start_epoch = load_checkpoint(
+            model, optimizer, scheduler, cfg.training.resume_checkpoint, cfg, accelerator
+        )
+        start_epoch += 1
+
+    # === Training Loop ===
+    log_message("Starting training...", cfg, accelerator)
+
+    best_val_loss = float("inf")
+
+    try:
+        for epoch in range(start_epoch, cfg.training.epochs + 1):
+            log_message(f"\n{'=' * 60}", cfg, accelerator)
+            log_message(f"EPOCH {epoch}/{cfg.training.epochs}", cfg, accelerator)
+            log_message(f"{'=' * 60}", cfg, accelerator)
+
+            global_step, best_val_loss = train_epoch(
+                model=model,
+                dataloader=train_dataloader,
+                optimizer=optimizer,
+                scheduler=scheduler,
+                cfg=cfg,
+                epoch=epoch,
+                global_step=global_step,
+                accelerator=accelerator,
+                val_dataloader=val_dataloader,
+                best_val_loss=best_val_loss,
+            )
+
+            if cfg.logging.save_every_epoch:
+                save_checkpoint(
+                    model, optimizer, scheduler, global_step, epoch, cfg, accelerator
+                )
+
+    except KeyboardInterrupt:
+        log_message("Training interrupted by user", cfg, accelerator)
+        save_checkpoint(model, optimizer, scheduler, global_step, epoch, cfg, accelerator)
+
+    # === Final Save ===
+    log_message("\nTraining completed!", cfg, accelerator)
+
+    if accelerator.is_main_process:
+        final_model_path = Path(cfg.paths.output_dir) / "model_final.pt"
+        unwrapped_model = accelerator.unwrap_model(model)
+        torch.save(unwrapped_model.state_dict(), final_model_path)
+        log_message(f"Final model: {final_model_path}", cfg, accelerator)
+
+    accelerator.wait_for_everyone()
+    finish_tracking()
+
+
+if __name__ == "__main__":
+    main()

pythia1b_v5_04_21/wandb/run-20260421_202839-8ing6xdi/run-8ing6xdi.wandb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:717527847fea27ac89ab840fa450ede0488f79e543e77544bf788b7b7673ba98
+size 5275648