openllava 3.0.0__py3-none-any.whl

openllava/__init__.py

@@ -0,0 +1,96 @@
+"""
+OpenLLaVA — Open-Source Multimodal Vision Injection Framework.
+
+Inject vision into any language model. Architecture-agnostic, multi-backend.
+
+Usage:
+    from openllava import OpenLLaVA, Backend, experts
+
+    model = OpenLLaVA(
+        llm="meta-llama/Llama-3-8B",
+        vision_encoder="google/siglip2-so400m-patch14-384",
+        backend=Backend.CUDA,
+    )
+    model.lora(r=64, alpha=128)
+    model.train(phase1=dict(dataset="liuhaotian/LLaVA-Pretrain", samples=100_000))
+    model.push("my-org/my-model")
+"""
+
+from __future__ import annotations
+
+from . import (
+    api,
+    backends,
+    cli,
+    data,
+    distributed,
+    eval,
+    experts,
+    export,
+    inference,
+    kernels,
+    optimizations,
+    rl,
+    serve,
+    training,
+    utils,
+)
+from .api import FastLanguageModel, FastVisionModel, OpenLLaVATrainer, TrainingConfig
+from .core.backend import Backend, BackendManager, get_backend, is_cuda_available
+from .core.model import OpenLLaVA
+from .core.patcher import AnyResProcessor, ModelPatcher, YakiModel, YakiProjector
+from .kernels import cuda_graphs
+from .kernels import triton as triton_kernels
+from .optimizations import (
+    EMAModel,
+    awq_quantize,
+    compile_model,
+    enable_fp8_training,
+    gptq_quantize,
+)
+from .utils import HardwareDetector, HardwareInfo, auto_configure, profile_model
+
+__version__ = "3.0.0"
+__author__ = "OpceanAI Research Team"
+
+__all__ = [
+    "OpenLLaVA",
+    "Backend",
+    "BackendManager",
+    "get_backend",
+    "is_cuda_available",
+    "YakiProjector",
+    "YakiModel",
+    "AnyResProcessor",
+    "ModelPatcher",
+    "experts",
+    "training",
+    "data",
+    "rl",
+    "export",
+    "eval",
+    "kernels",
+    "optimizations",
+    "inference",
+    "distributed",
+    "serve",
+    "backends",
+    "api",
+    "cli",
+    "utils",
+    "triton_kernels",
+    "cuda_graphs",
+    "compile_model",
+    "enable_fp8_training",
+    "EMAModel",
+    "gptq_quantize",
+    "awq_quantize",
+    "FastVisionModel",
+    "FastLanguageModel",
+    "OpenLLaVATrainer",
+    "TrainingConfig",
+    "auto_configure",
+    "HardwareDetector",
+    "HardwareInfo",
+    "profile_model",
+]

openllava/api/__init__.py

@@ -0,0 +1,31 @@
+"""OpenLLaVA High-Level API — Unsloth-compatible fast model loading and training.
+
+Provides drop-in replacements for common training frameworks with
+auto-configuration, multi-backend logging, and callback orchestration.
+
+Usage:
+    from openllava.api import FastVisionModel, OpenLLaVATrainer, TrainingConfig
+
+    model, tokenizer = FastVisionModel.from_pretrained("openllava/yaki-8b")
+    config = TrainingConfig(mode="lora", output_dir="./output")
+    trainer = OpenLLaVATrainer(model=model, tokenizer=tokenizer, args=config)
+    trainer.train()
+    trainer.save_model("./final")
+"""
+
+from __future__ import annotations
+
+from .config import TrainingConfig
+from .fast_model import FastLanguageModel, FastVisionModel
+from .strategies import auto_configure, get_peft_model, load_dataset
+from .trainer import OpenLLaVATrainer
+
+__all__ = [
+    "FastVisionModel",
+    "FastLanguageModel",
+    "OpenLLaVATrainer",
+    "TrainingConfig",
+    "get_peft_model",
+    "load_dataset",
+    "auto_configure",
+]

openllava/api/callbacks.py

@@ -0,0 +1,463 @@
+"""OpenLLaVA Callback System.
+
+Complete training callback infrastructure with hooks for monitoring, checkpointing,
+early stopping, and memory profiling. Designed for composition via CallbackList.
+
+Usage:
+    from openllava.api.callbacks import EarlyStoppingCallback, CallbackList
+
+    callbacks = CallbackList([
+        EarlyStoppingCallback(monitor="loss", patience=3),
+        ModelCheckpointCallback(dirpath="./checkpoints", monitor="loss"),
+    ])
+    callbacks.on_train_begin()
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import time
+from abc import ABC
+from typing import Any
+
+import torch
+
+_HOOKS = [
+    "on_train_begin", "on_train_end",
+    "on_epoch_begin", "on_epoch_end",
+    "on_step_begin", "on_step_end",
+    "on_save", "on_log", "on_evaluate",
+]
+
+
+class BaseCallback(ABC):
+    """Abstract base for all training callbacks.
+
+    Subclasses override any combination of the hook methods below.
+    All hooks receive **kwargs for forward compatibility.
+    """
+
+    def on_train_begin(self, **kwargs) -> None:
+        pass
+
+    def on_train_end(self, **kwargs) -> None:
+        pass
+
+    def on_epoch_begin(self, **kwargs) -> None:
+        pass
+
+    def on_epoch_end(self, **kwargs) -> None:
+        pass
+
+    def on_step_begin(self, **kwargs) -> None:
+        pass
+
+    def on_step_end(self, **kwargs) -> None:
+        pass
+
+    def on_save(self, **kwargs) -> None:
+        pass
+
+    def on_log(self, **kwargs) -> None:
+        pass
+
+    def on_evaluate(self, **kwargs) -> None:
+        pass
+
+
+class CallbackList:
+    """Manages a list of callbacks and dispatches hooks to all of them.
+
+    Usage:
+        cbl = CallbackList([EarlyStoppingCallback(), ModelCheckpointCallback()])
+        cbl.on_train_begin(trainer=trainer)
+    """
+
+    def __init__(self, callbacks: list[BaseCallback] | None = None):
+        self._callbacks: list[BaseCallback] = list(callbacks) if callbacks else []
+
+    def append(self, callback: BaseCallback) -> None:
+        self._callbacks.append(callback)
+
+    def extend(self, callbacks: list[BaseCallback]) -> None:
+        self._callbacks.extend(callbacks)
+
+    def dispatch(self, hook: str, **kwargs) -> None:
+        for cb in self._callbacks:
+            try:
+                getattr(cb, hook)(**kwargs)
+            except Exception as e:
+                print(f"[CallbackList] Error in {cb.__class__.__name__}.{hook}: {e}")
+
+    def on_train_begin(self, **kwargs) -> None:
+        self.dispatch("on_train_begin", **kwargs)
+
+    def on_train_end(self, **kwargs) -> None:
+        self.dispatch("on_train_end", **kwargs)
+
+    def on_epoch_begin(self, **kwargs) -> None:
+        self.dispatch("on_epoch_begin", **kwargs)
+
+    def on_epoch_end(self, **kwargs) -> None:
+        self.dispatch("on_epoch_end", **kwargs)
+
+    def on_step_begin(self, **kwargs) -> None:
+        self.dispatch("on_step_begin", **kwargs)
+
+    def on_step_end(self, **kwargs) -> None:
+        self.dispatch("on_step_end", **kwargs)
+
+    def on_save(self, **kwargs) -> None:
+        self.dispatch("on_save", **kwargs)
+
+    def on_log(self, **kwargs) -> None:
+        self.dispatch("on_log", **kwargs)
+
+    def on_evaluate(self, **kwargs) -> None:
+        self.dispatch("on_evaluate", **kwargs)
+
+    def __len__(self) -> int:
+        return len(self._callbacks)
+
+    def __iter__(self):
+        return iter(self._callbacks)
+
+
+class EarlyStoppingCallback(BaseCallback):
+    """Stop training when a monitored metric stops improving.
+
+    Args:
+        monitor: Metric name to monitor (default: "loss").
+        patience: Number of validation checks with no improvement before stopping.
+        min_delta: Minimum change to qualify as an improvement.
+        mode: "min" (lower is better) or "max" (higher is better).
+    """
+
+    def __init__(
+        self,
+        monitor: str = "loss",
+        patience: int = 3,
+        min_delta: float = 0.0,
+        mode: str = "min",
+    ):
+        if patience < 1:
+            raise ValueError(f"patience must be >= 1, got {patience}")
+        if mode not in {"min", "max"}:
+            raise ValueError(f"mode must be 'min' or 'max', got '{mode}'")
+        self.monitor = monitor
+        self.patience = patience
+        self.min_delta = min_delta
+        self.mode = mode
+        self._best: float | None = None
+        self._counter: int = 0
+        self._should_stop: bool = False
+
+    def _is_improvement(self, current: float) -> bool:
+        if self._best is None:
+            return True
+        if self.mode == "min":
+            return current < self._best - self.min_delta
+        return current > self._best + self.min_delta
+
+    def on_train_begin(self, **kwargs) -> None:
+        self._best = None
+        self._counter = 0
+        self._should_stop = False
+
+    def on_evaluate(self, **kwargs) -> None:
+        current = kwargs.get("metrics", {}).get(self.monitor, None)
+        if current is None:
+            return
+        if self._is_improvement(current):
+            self._best = current
+            self._counter = 0
+        else:
+            self._counter += 1
+            if self._counter >= self.patience:
+                self._should_stop = True
+                print(f"[EarlyStopping] Stopping after {self._counter} "
+                      f"checks without improvement in {self.monitor}")
+
+    @property
+    def should_stop(self) -> bool:
+        return self._should_stop
+
+    @property
+    def best_value(self) -> float | None:
+        return self._best
+
+
+class ModelCheckpointCallback(BaseCallback):
+    """Save model checkpoints during training.
+
+    Args:
+        dirpath: Directory to save checkpoints.
+        monitor: Metric to monitor for best model selection.
+        save_top_k: Number of best checkpoints to keep (-1 for all).
+        mode: "min" or "max" for metric comparison.
+        save_last: Whether to save the last checkpoint.
+    """
+
+    def __init__(
+        self,
+        dirpath: str = "./checkpoints",
+        monitor: str = "loss",
+        save_top_k: int = 1,
+        mode: str = "min",
+        save_last: bool = True,
+    ):
+        self.dirpath = dirpath
+        self.monitor = monitor
+        self.save_top_k = save_top_k
+        self.mode = mode
+        self.save_last = save_last
+        self._best_scores: list[tuple[float, str]] = []
+        os.makedirs(dirpath, exist_ok=True)
+
+    def on_evaluate(self, **kwargs) -> None:
+        score = kwargs.get("metrics", {}).get(self.monitor)
+        if score is None:
+            return
+        step = kwargs.get("step", 0)
+        model = kwargs.get("model")
+        if model is None:
+            return
+
+        path = os.path.join(self.dirpath, f"checkpoint-{step}")
+        self._save_model(model, path)
+
+        if self.save_top_k > 0:
+            self._best_scores.append((float(score), path))
+            self._best_scores.sort(
+                key=lambda x: x[0],
+                reverse=(self.mode == "max"),
+            )
+            while len(self._best_scores) > self.save_top_k:
+                _, old_path = self._best_scores.pop()
+                self._remove_checkpoint(old_path)
+
+    def on_train_end(self, **kwargs) -> None:
+        if self.save_last:
+            model = kwargs.get("model")
+            if model is not None:
+                path = os.path.join(self.dirpath, "last")
+                self._save_model(model, path)
+
+    def _save_model(self, model, path: str) -> None:
+        try:
+            os.makedirs(path, exist_ok=True)
+            if hasattr(model, "save_pretrained"):
+                model.save_pretrained(path)
+            else:
+                torch.save(model.state_dict(), os.path.join(path, "pytorch_model.bin"))
+        except Exception as e:
+            print(f"[ModelCheckpoint] Failed to save to {path}: {e}")
+
+    @staticmethod
+    def _remove_checkpoint(path: str) -> None:
+        import shutil
+        if os.path.exists(path):
+            shutil.rmtree(path, ignore_errors=True)
+
+
+class GradientAccumulationCallback(BaseCallback):
+    """Dynamic gradient accumulation step adjustment.
+
+    Adjusts accumulation steps based on observed batch memory usage
+    to fit within available GPU memory.
+    """
+
+    def __init__(self, target_effective_batch: int = 32, min_accum: int = 1, max_accum: int = 128):
+        self.target = target_effective_batch
+        self.min_accum = min_accum
+        self.max_accum = max_accum
+        self.current_accum: int = 1
+        self._oom_count: int = 0
+
+    def on_train_begin(self, **kwargs) -> None:
+        self.current_accum = getattr(
+            kwargs.get("trainer", {}), "gradient_accumulation_steps", 1
+        )
+
+    def on_step_end(self, **kwargs) -> None:
+        loss = kwargs.get("loss", None)
+        if loss is not None and torch.isnan(loss):
+            self.current_accum = min(self.current_accum * 2, self.max_accum)
+            self._oom_count += 1
+            print(f"[GradientAccum] NaN loss detected, increasing accumulation to {self.current_accum}")
+
+
+class LearningRateMonitorCallback(BaseCallback):
+    """Log learning rate at each step."""
+
+    def __init__(self):
+        self._lrs: list[tuple[int, float]] = []
+
+    def on_step_end(self, **kwargs) -> None:
+        step = kwargs.get("step", len(self._lrs))
+        optimizer = kwargs.get("optimizer")
+        if optimizer is not None and len(optimizer.param_groups) > 0:
+            lr = optimizer.param_groups[0]["lr"]
+            self._lrs.append((step, lr))
+
+    def get_history(self) -> list[tuple[int, float]]:
+        return self._lrs
+
+    def on_train_end(self, **kwargs) -> None:
+        self._lrs.clear()
+
+
+class MemoryProfilerCallback(BaseCallback):
+    """Log GPU/CPU memory usage periodically."""
+
+    def __init__(self, log_every_n_steps: int = 100):
+        self.log_every = log_every_n_steps
+        self._step: int = 0
+        self._peak_gpu_mb: float = 0.0
+
+    def on_step_end(self, **kwargs) -> None:
+        self._step += 1
+        if self._step % self.log_every != 0:
+            return
+        info_parts = []
+        if torch.cuda.is_available():
+            try:
+                import pynvml
+                pynvml.nvmlInit()
+                for i in range(torch.cuda.device_count()):
+                    handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+                    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                    gb = info.used / (1024 ** 3)
+                    self._peak_gpu_mb = max(self._peak_gpu_mb, info.used / (1024 ** 2))
+                    info_parts.append(f"GPU{i}: {gb:.1f}GB")
+            except ImportError:
+                alloc = torch.cuda.memory_allocated() / (1024 ** 3)
+                reserved = torch.cuda.memory_reserved() / (1024 ** 3)
+                info_parts.append(f"GPU alloc: {alloc:.1f}GB, reserved: {reserved:.1f}GB")
+        import psutil
+        try:
+            ram = psutil.virtual_memory()
+            info_parts.append(f"RAM: {ram.percent:.0f}%")
+        except ImportError:
+            pass
+        if info_parts:
+            print(f"[Memory] Step {self._step}: {' | '.join(info_parts)}")
+
+
+class TimingCallback(BaseCallback):
+    """Track epoch and batch timing statistics."""
+
+    def __init__(self):
+        self._epoch_start: float = 0.0
+        self._step_start: float = 0.0
+        self.epoch_times: list[float] = []
+        self.step_times: list[float] = []
+
+    def on_train_begin(self, **kwargs) -> None:
+        self.epoch_times.clear()
+        self.step_times.clear()
+
+    def on_epoch_begin(self, **kwargs) -> None:
+        self._epoch_start = time.perf_counter()
+
+    def on_epoch_end(self, **kwargs) -> None:
+        elapsed = time.perf_counter() - self._epoch_start
+        self.epoch_times.append(elapsed)
+
+    def on_step_begin(self, **kwargs) -> None:
+        self._step_start = time.perf_counter()
+
+    def on_step_end(self, **kwargs) -> None:
+        elapsed = time.perf_counter() - self._step_start
+        self.step_times.append(elapsed)
+
+    def summary(self) -> dict[str, float]:
+        if not self.step_times:
+            return {}
+        return {
+            "mean_step_s": sum(self.step_times) / len(self.step_times),
+            "min_step_s": min(self.step_times),
+            "max_step_s": max(self.step_times),
+            "mean_epoch_s": sum(self.epoch_times) / len(self.epoch_times) if self.epoch_times else 0.0,
+        }
+
+
+class ProgressBarCallback(BaseCallback):
+    """Rich progress bar display for training progress."""
+
+    def __init__(self, total_steps: int = 0, use_rich: bool = False):
+        self.total_steps = total_steps
+        self._current: int = 0
+        self._pbar: Any = None
+        self._use_rich = use_rich
+
+    def on_train_begin(self, **kwargs) -> None:
+        if self._use_rich:
+            try:
+                from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
+                self._rich_progress = Progress(
+                    TextColumn("[progress.description]{task.description}"),
+                    BarColumn(),
+                    TextColumn("{task.completed}/{task.total}"),
+                    TimeRemainingColumn(),
+                )
+                self._rich_progress.start()
+                self._task = self._rich_progress.add_task("[cyan]Training...", total=self.total_steps)
+            except ImportError:
+                self._use_rich = False
+        if not self._use_rich:
+            from tqdm import tqdm
+            self._pbar = tqdm(total=self.total_steps, desc="Training", unit="step")
+
+    def on_step_end(self, **kwargs) -> None:
+        self._current += 1
+        loss = kwargs.get("loss", None)
+        postfix = {}
+        if loss is not None:
+            postfix["loss"] = f"{loss:.4f}" if isinstance(loss, (int, float)) else str(loss)
+        if self._use_rich and hasattr(self, "_rich_progress"):
+            self._rich_progress.update(self._task, advance=1)
+        elif self._pbar is not None:
+            if postfix:
+                self._pbar.set_postfix(postfix)
+            self._pbar.update(1)
+
+    def on_train_end(self, **kwargs) -> None:
+        if self._use_rich and hasattr(self, "_rich_progress"):
+            self._rich_progress.stop()
+        elif self._pbar is not None:
+            self._pbar.close()
+
+
+class NaNMonitorCallback(BaseCallback):
+    """Stop training if NaN values are detected in loss or gradients."""
+
+    def __init__(self, max_nan_count: int = 5):
+        self.max_nan_count = max_nan_count
+        self._nan_count: int = 0
+        self._should_stop: bool = False
+
+    def on_train_begin(self, **kwargs) -> None:
+        self._nan_count = 0
+        self._should_stop = False
+
+    def on_step_end(self, **kwargs) -> None:
+        loss = kwargs.get("loss", None)
+        if loss is None:
+            return
+        loss_val = loss.item() if hasattr(loss, "item") else float(loss)
+        grad_norm = kwargs.get("grad_norm", None)
+        is_nan = math.isnan(loss_val) or math.isinf(loss_val)
+        if grad_norm is not None and hasattr(grad_norm, "item"):
+            gn = grad_norm.item() if hasattr(grad_norm, "item") else float(grad_norm)
+            is_nan = is_nan or math.isnan(gn) or math.isinf(gn)
+        if is_nan:
+            self._nan_count += 1
+            print(f"[NaNMonitor] NaN/Inf detected ({self._nan_count}/{self.max_nan_count})")
+            if self._nan_count >= self.max_nan_count:
+                self._should_stop = True
+
+    @property
+    def should_stop(self) -> bool:
+        return self._should_stop