langtune 0.1.19__py3-none-any.whl

langtune/model/hub.py

@@ -0,0 +1,109 @@
+"""
+Hub Resolver for efficient model downloads.
+
+Handles interactions with the Hugging Face Hub, including:
+- Snapshot downloads with caching
+- Authentication
+- Offline mode support
+- File resolution
+"""
+
+import os
+import logging
+from typing import Optional, List, Dict, Union, Path
+from pathlib import Path as PathLib
+from huggingface_hub import snapshot_download, get_token, HfFolder
+from huggingface_hub.utils import LocalEntryNotFoundError, EntryNotFoundError, RevisionNotFoundError, RepositoryNotFoundError
+
+logger = logging.getLogger(__name__)
+
+class HubResolver:
+    """
+    Resolves and downloads model files from the Hugging Face Hub.
+    
+    Optimized for:
+    - Speed (concurrent downloads)
+    - Caching (avoid re-downloading)
+    - Offline usage (finding local files)
+    """
+    
+    def __init__(
+        self,
+        cache_dir: Optional[str] = None,
+        local_files_only: bool = False,
+        token: Optional[str] = None,
+    ):
+        self.cache_dir = cache_dir or os.environ.get("LANGTUNE_CACHE_DIR")
+        self.local_files_only = local_files_only
+        self.token = token or get_token()
+        
+    def resolve(
+        self,
+        model_id: str,
+        revision: str = "main",
+        allow_patterns: Optional[List[str]] = None,
+        ignore_patterns: Optional[List[str]] = None,
+    ) -> PathLib:
+        """
+        Download or find model snapshot.
+        
+        Args:
+            model_id: The model ID on HF Hub (e.g. 'meta-llama/Llama-2-7b-hf')
+            revision: Branch or commit hash
+            allow_patterns: Files to include (default: ['*.safetensors', '*.json', '*.model'])
+            ignore_patterns: Files to exclude (default: ['*.bin', '*.pth'])
+            
+        Returns:
+            Path to the directory containing the model files.
+        """
+        # Default patterns to prioritize safetensors and configs
+        if allow_patterns is None:
+            allow_patterns = ["*.safetensors", "*.json", "*.model", "tokenizer*"]
+            
+        if ignore_patterns is None:
+            # Explicitly ignore pytorch/pickle weights if we want to enforce safetensors
+            ignore_patterns = ["*.bin", "*.pth", "*.pt"]
+
+        logger.info(f"Resolving model {model_id} (revision={revision})...")
+        
+        try:
+            download_path = snapshot_download(
+                repo_id=model_id,
+                revision=revision,
+                cache_dir=self.cache_dir,
+                local_files_only=self.local_files_only,
+                token=self.token,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                resume_download=True,
+                max_workers=8, # Parallel downloads
+                tqdm_class=None, # We might want to hook a custom progress bar later
+            )
+            
+            logger.info(f"Model resolved to: {download_path}")
+            return PathLib(download_path)
+            
+        except (LocalEntryNotFoundError, EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError) as e:
+            if self.local_files_only:
+                logger.error(f"Model {model_id} not found locally and offline mode is enabled.")
+                raise RuntimeError(f"Model {model_id} not found locally. Disable offline mode to download.") from e
+            else:
+                logger.error(f"Failed to download model {model_id}: {e}")
+                raise RuntimeError(f"Failed to resolve model {model_id}: {e}") from e
+                
+        except Exception as e:
+            logger.error(f"Unexpected error resolving model {model_id}: {e}")
+            raise
+
+    @property
+    def is_offline(self) -> bool:
+        return self.local_files_only
+
+# Global instance capabilities
+_default_resolver = None
+
+def get_resolver(offline: bool = False, cache_dir: Optional[str] = None) -> HubResolver:
+    global _default_resolver
+    if _default_resolver is None or _default_resolver.local_files_only != offline:
+        _default_resolver = HubResolver(local_files_only=offline, cache_dir=cache_dir)
+    return _default_resolver

langtune/model/loader.py

@@ -0,0 +1,84 @@
+"""
+Model Loader.
+
+Orchestrates the loading pipeline:
+1. Hub Resolution
+2. Architecture Construction
+3. Streaming Weight Loading
+4. Kernel Injection
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Dict, Any, Union
+import logging
+
+from .hub import HubResolver, get_resolver
+from .safetensors import TensorStreamer
+from .weights import WeightLoader
+
+logger = logging.getLogger(__name__)
+
+class ModelLoader:
+    """
+    Main entry point for loading models with high performance.
+    """
+    
+    def __init__(
+        self, 
+        offline: bool = False,
+        cache_dir: Optional[str] = None
+    ):
+        self.resolver = get_resolver(offline=offline, cache_dir=cache_dir)
+        
+    def load(
+        self,
+        model_id: str,
+        quantization: Optional[str] = None, # "nf4", "fp4"
+        dtype: str = "bf16",
+        device: str = "cuda"
+    ) -> nn.Module:
+        """
+        Load a model with the optimized pipeline.
+        """
+        # 1. Resolve
+        model_path = self.resolver.resolve(model_id)
+        
+        # 2. Config & Architecture
+        # For now, we rely on HF config to build the skeleton
+        # In the future, we will use our "architectures" registry
+        from transformers import AutoConfig, AutoModelForCausalLM
+        
+        config = AutoConfig.from_pretrained(model_path)
+        
+        logger.info(f"Building model skeleton for {model_id}...")
+        # Creation on 'meta' device avoids memory allocation!
+        with torch.device("meta"):
+             model = AutoModelForCausalLM.from_config(config)
+             
+        # Move empty shell to CPU/device?
+        # Meta tensors can't be used directly for loading usually without `to_empty`
+        # Using `to_empty` moves to device but allocating memory
+        model = model.to_empty(device=device)
+        
+        # 3. Stream Weights
+        logger.info("Streaming weights...")
+        streamer = TensorStreamer(model_path)
+        
+        # 4. Load & Quantize
+        compute_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16
+        loader = WeightLoader(streamer, quantization_mode=quantization, compute_dtype=compute_dtype, device=device)
+        
+        loader.load_into_module(model)
+        
+        # 5. Kernel Injection (Placeholder)
+        # self._inject_kernels(model)
+        
+        return model
+
+    def _inject_kernels(self, model: nn.Module):
+        """
+        Replace layers with Langtrain Custom Kernels.
+        """
+        # Implementation to be added in next step
+        pass

langtune/model/safetensors.py

@@ -0,0 +1,104 @@
+"""
+Safetensors Streamer.
+
+Efficiently streams tensors from disk using memory mapping.
+"""
+
+import os
+import logging
+import torch
+from typing import Dict, List, Optional, Union, Iterator, Tuple
+from pathlib import Path
+from safetensors import safe_open
+from dataclasses import dataclass
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class TensorInfo:
+    """Metadata about a tensor in a safetensors file."""
+    name: str
+    shape: List[int]
+    dtype: str
+    file_path: str
+
+class TensorStreamer:
+    """
+    Manages lazy loading of tensors from a directory of safetensors files.
+    """
+    
+    def __init__(self, model_dir: Union[str, Path]):
+        self.model_dir = Path(model_dir)
+        self.files = sorted(list(self.model_dir.glob("*.safetensors")))
+        
+        if not self.files:
+            raise FileNotFoundError(f"No .safetensors files found in {self.model_dir}")
+            
+        self.index: Dict[str, TensorInfo] = {}
+        self._build_index()
+        
+    def _build_index(self):
+        """Scan files and build an index of tensor locations."""
+        logger.info(f"Indexing {len(self.files)} safetensors files...")
+        
+        for file_path in self.files:
+            try:
+                with safe_open(file_path, framework="pt", device="cpu") as f:
+                    for key in f.keys():
+                        # We don't load the tensor, just inspecting keys
+                        # Note: safe_open doesn't give shape/dtype without loading or accessing metadata
+                        # Ideally we read the header. safe_open provides proper access.
+                        # For pure metadata scan without loading payload, external tools or specialized logic is best.
+                        # But standard safe_open is memory efficient (mmap).
+                        
+                        # We store just the file mapping for now to keep it fast.
+                        # Conflict check
+                        if key in self.index:
+                            logger.warning(f"Duplicate tensor {key} found in {file_path}, ignoring duplicate.")
+                            continue
+                            
+                        self.index[key] = TensorInfo(
+                            name=key,
+                            shape=[], # Lazy populate if needed, or read tensor for shape cheap
+                            dtype="",
+                            file_path=str(file_path)
+                        )
+            except Exception as e:
+                logger.error(f"Failed to read {file_path}: {e}")
+                raise
+                
+    def has_tensor(self, name: str) -> bool:
+        return name in self.index
+        
+    def get_tensor(self, name: str, device: str = "cpu") -> torch.Tensor:
+        """
+        Load a single tensor.
+        """
+        if name not in self.index:
+            raise KeyError(f"Tensor {name} not found in model files.")
+            
+        info = self.index[name]
+        
+        with safe_open(info.file_path, framework="pt", device=device) as f:
+            return f.get_tensor(name)
+
+    def load_state_dict(self, device: str = "cpu") -> Dict[str, torch.Tensor]:
+        """
+        Load all tensors into a state dict (WARNING: High Memory Usage).
+        Use for small models or debugging only.
+        """
+        state_dict = {}
+        for name in self.index:
+            state_dict[name] = self.get_tensor(name, device)
+        return state_dict
+
+    def stream(self, device: str = "cpu") -> Iterator[Tuple[str, torch.Tensor]]:
+        """
+        Yields (name, tensor) pairs one by one to minimize memory usage.
+        """
+        # Optimize by opening files once
+        for file_path in self.files:
+            with safe_open(file_path, framework="pt", device=device) as f:
+                for key in f.keys():
+                    yield key, f.get_tensor(key)
+

langtune/model/weights.py

@@ -0,0 +1,100 @@
+"""
+Weight Loader.
+
+Handles reading weights from the streamer and applying:
+- Quantize-on-Read (NF4/FP4)
+- Precision casting (BF16/FP16)
+- Kernel weight format definition
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict, Any, Optional
+import logging
+
+from .safetensors import TensorStreamer
+
+logger = logging.getLogger(__name__)
+
+class WeightLoader:
+    """
+    Loads weights into a model structure.
+    
+    Supports:
+    - Direct quantization (NF4) during loading to save memory
+    - BF16/FP16 casting
+    - Kernel-specific formatting
+    """
+    
+    def __init__(
+        self,
+        streamer: TensorStreamer,
+        quantization_mode: Optional[str] = None, # "nf4", "fp4", "int8"
+        compute_dtype: torch.dtype = torch.bfloat16,
+        device: str = "cuda"
+    ):
+        self.streamer = streamer
+        self.quantization_mode = quantization_mode
+        self.compute_dtype = compute_dtype
+        self.device = device
+        
+    def load_into_module(self, module: nn.Module, prefix: str = ""):
+        """
+        Populate a module's weights from the streamer.
+        """
+        # Get all parameters/buffers in the module
+        # state_dict keys are relative to the module if we use module.named_parameters()
+        # But streamer has full keys (e.g. "model.layers.0.self_attn.q_proj.weight")
+        
+        # We assume the model structure matches the streamer keys for now
+        # Ideally, we iterate over the model parameters and fetch corresponding keys
+        
+        for name, param in module.named_parameters(recurse=False):
+            full_name = f"{prefix}.{name}" if prefix else name
+            
+            if self.streamer.has_tensor(full_name):
+                tensor = self.streamer.get_tensor(full_name, device="cpu")
+                
+                # Apply Quantize-on-Read
+                if self.quantization_mode and "lora" not in full_name and "norm" not in full_name:
+                     self._quantize_and_assign(module, name, tensor)
+                else:
+                    # Standard loading
+                    with torch.no_grad():
+                        param.data = tensor.to(self.device, dtype=self.compute_dtype)
+                        
+        # Recursively load children
+        for child_name, child in module.named_children():
+            child_prefix = f"{prefix}.{child_name}" if prefix else child_name
+            self.load_into_module(child, prefix=child_prefix)
+
+    def _quantize_and_assign(self, module: nn.Module, param_name: str, tensor: torch.Tensor):
+        """
+        Quantize tensor and replace parameter in module.
+        Currently falls back to bitsandbytes functional quantization if usage is detected,
+        or custom kernel logic.
+        """
+        # Placeholder for NF4 quantization connection
+        # In a real implementation, we would use bitsandbytes.functional.quantize_4bit
+        # and replace the nn.Parameter with a Params4bit
+        
+        if self.quantization_mode == "nf4":
+             try:
+                 import bitsandbytes as bnb
+                 from bitsandbytes.nn import Params4bit
+                 
+                 # Quantize to NF4
+                 # This requires GPU generally
+                 tensor_gpu = tensor.to(self.device)
+                 param_4bit = Params4bit(tensor_gpu, requires_grad=False, compress_statistics=True, quant_type="nf4")
+                 
+                 # Replace parameter
+                 setattr(module, param_name, param_4bit)
+                 
+             except ImportError:
+                 logger.warning("bitsandbytes not found, falling back to BF16/FP16")
+                 with torch.no_grad():
+                     getattr(module, param_name).data = tensor.to(self.device, dtype=self.compute_dtype)
+        else:
+             with torch.no_grad():
+                getattr(module, param_name).data = tensor.to(self.device, dtype=self.compute_dtype)

langtune/models.py

@@ -0,0 +1,19 @@
+"""
+models.py: LoRA-enabled transformer models for Langtune.
+
+This module now re-exports components from `langtune.nn` for backward compatibility.
+"""
+
+from .nn.layers import LoRALinear, MultiHeadAttention
+from .nn.transformer import TransformerBlock, LoRALanguageModel
+from .nn.fast_transformer import FastMultiHeadAttention, FastTransformerBlock, FastLoRALanguageModel
+
+__all__ = [
+    "LoRALinear",
+    "MultiHeadAttention",
+    "TransformerBlock",
+    "LoRALanguageModel",
+    "FastMultiHeadAttention",
+    "FastTransformerBlock",
+    "FastLoRALanguageModel",
+]