Upload folder using huggingface_hub

ztrain/io.py

@@ -0,0 +1,30 @@
+# ztrain/io.py
+# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted
+
+import os
+from glob import glob
+
+def flatten_index(model_paths : list[str], allow_list : list[str]):
+    flat = []
+    subtype = []
+    index = {}
+    ix = 0
+    for g in sorted(model_paths):
+        name = os.path.basename(g)
+        if name in allow_list:
+            index[name] = ix
+            flat.append(name)
+            if 'base' in g:
+                subtype.append('base')
+            elif 'instruct' in g:
+                subtype.append('instruct')
+            else:
+                subtype.append('other')
+        ix += 1
+    return index, flat, subtype
+
+def list_for_path(path: str, include_folders: list[str], search: str = "/**/*") -> tuple[list[str], list[str], list[str], dict[str, int]]:
+    model_list = sorted([*[ f for f in glob(path + search)]])
+    group_idx, model_names, subtypes = flatten_index(model_list, include_folders)
+    groups = [[m] for m in model_names]
+    return model_names, subtypes, model_list, group_idx

ztrain/model.py

@@ -0,0 +1,39 @@
+# ztrain/model.py
+# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted
+
+from collections import defaultdict
+import re
+
+def generate_merge_group(group_data : list, parents : list[int] = []):
+    # drill down until we find a list of strings, then yield it with a parent tree index
+    for i, g in enumerate(group_data):
+        if isinstance(g, list):
+            yield from generate_merge_group(g, parents + [i])
+        else:
+            yield g, parents + [i]
+
+def merge_groups(group_data : list):
+    results = defaultdict(list)
+    for g, k in generate_merge_group(group_data):
+        key = tuple(k[:-1])
+        results[key].append(g)
+    return results
+
+def get_layer_type(k : str) -> tuple[int, str, str, str]:
+    matcher = re.compile(r"model.layers.(\d+)\.(.+)\.(.+)\.(.+)")
+
+    m = matcher.match(k)
+    if m is not None:
+        return int(m.group(1)), m.group(2), m.group(3), m.group(4)
+    matcher = re.compile(r"model.layers.(\d+)\.(.+)\.(.+)")
+    if m is not None:
+        return int(m.group(1)), m.group(2), "", m.group(3)
+
+    if "model.norm.weight" == k:
+        return -1, "norm", "", "weight"
+    if "model.embed_tokens.weight" == k:
+        return -1, "embed_tokens", "", "weight"
+    if "lm_head.weight" == k:
+        return -1, "lm_head", "", "weight"
+    print(f"Unknown key {k}")
+    return -1, "unknown", "unknown", "unknown"

ztrain/signal.py

@@ -0,0 +1,79 @@
+# ztrain/signal.py
+# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted
+
+import torch
+
+def gaussian_kernel(size, sigma=1.0):
+    """
+    Generates a 2D Gaussian kernel using PyTorch.
+
+    Parameters:
+    - size: The size of the kernel (an integer). It's recommended to use an odd number
+            to have a central pixel.
+    - sigma: The standard deviation of the Gaussian distribution.
+
+    Returns:
+    - A 2D PyTorch tensor representing the Gaussian kernel.
+    """
+    size = int(size) // 2
+    x, y = torch.meshgrid(torch.arange(-size, size+1), torch.arange(-size, size+1))
+    g = torch.exp(-(x**2 + y**2) / (2 * sigma**2))
+    return g / g.sum()
+
+def laplacian_kernel(size, scale=1.0):
+    """
+    Creates a Laplacian kernel for edge detection with an adjustable size and scale factor.
+
+    Parameters:
+    - size: The size of the kernel (an integer). It's recommended to use an odd number
+            to ensure a central pixel.
+    - scale: A float that adjusts the intensity of the edge detection effect.
+
+    Returns:
+    - A 2D PyTorch tensor representing the scaled Laplacian kernel.
+    """
+    if size % 2 == 0:
+        raise ValueError("Size must be odd.")
+    
+    # Initialize the kernel with zeros
+    kernel = torch.zeros((size, size), dtype=torch.float32)
+    
+    # Set the center pixel
+    kernel[size // 2, size // 2] = -4.0
+    
+    # Set the immediate neighbors
+    kernel[size // 2, size // 2 - 1] = kernel[size // 2, size // 2 + 1] = 1.0
+    kernel[size // 2 - 1, size // 2] = kernel[size // 2 + 1, size // 2] = 1.0
+    
+    # For larger kernels, adjust the outer pixels (this simplistic approach might need refinement for larger sizes)
+    if size > 3:
+        for i in range(size):
+            for j in range(size):
+                if i == 0 or i == size - 1 or j == 0 or j == size - 1:
+                    kernel[i, j] = 1.0
+
+    # Apply the scale factor
+    kernel *= scale
+    
+    # Adjust the kernel so that its sum is 0
+    center = size // 2
+    kernel[center, center] = -torch.sum(kernel) + kernel[center, center]
+    
+    return kernel
+
+def fftshift(input):
+    """
+    Reorients the FFT output so the zero-frequency component is at the center.
+
+    Parameters:
+    - input: A 2D tensor representing the FFT output.
+
+    Returns:
+    - A 2D tensor with the zero-frequency component shifted to the center.
+    """
+    # For even dimensions, we split at dim_size // 2. For odd dimensions, we need to do (dim_size + 1) // 2
+    for dim in range(2):  # assuming input is 2D
+        n = input.shape[dim]
+        half = (n + 1) // 2
+        input = torch.roll(input, shifts=half, dims=dim)
+    return input

ztrain/stats.py

@@ -0,0 +1,30 @@
+# ztrain/stats.py
+# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted
+
+import os
+import torch
+from typing import Optional
+
+def gen_stats(delta : torch.Tensor, base : Optional[torch.Tensor]) -> tuple[float, float, float, float]:
+    if base is None:
+        rebuilt = delta
+    else:
+        rebuilt = base + delta
+    norm = rebuilt.norm().item()
+    if base is None:
+        cosine = 0
+    else:
+        cosine = torch.nn.functional.cosine_similarity(rebuilt, base, dim=0).mean().item()
+    min = delta.min().item()
+    max = delta.max().item()
+    del rebuilt
+    return norm, cosine, min, max
+
+def get_report(m0: torch.Tensor, stack : torch.Tensor, model_list : list[str]):
+    norm, cosine, min, max = gen_stats(m0, None)
+    print(f"Base Model {norm} {min} {max}")
+
+    for i, s in enumerate(stack):
+        model_name = os.path.basename(model_list[i])
+        norm, cosine, min, max = gen_stats(s, m0)
+        print(f"{model_name} {norm} {cosine} {min} {max}")

ztrain/tensors.py

@@ -0,0 +1,258 @@
+# ztrain/tensors.py
+# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted
+
+import torch
+from typing import Generator, Tuple
+
+def normalize_to(m1 : torch.Tensor, norm : torch.float32) -> tuple[torch.Tensor, torch.float32, torch.float32]:
+    m1 = m1.to(torch.float32)
+    m1_norm = torch.norm(m1)
+    ratio = (norm / m1_norm).item()
+    m1 = m1 * ratio
+    return m1, norm.item(), ratio
+
+def norm_ratio(m1 : torch.Tensor, m2 : torch.Tensor) -> float:
+    m1_norm = torch.norm(m1)
+    m2_norm = torch.norm(m2)
+    ratio = (m1_norm / m2_norm).item()
+    print(f"Norms {m1_norm} {m2_norm} {ratio}")
+    return ratio
+
+def merge_tensors_fft2(v0: torch.Tensor, v1: torch.Tensor, t: float) -> torch.Tensor:
+    """
+    Merges two tensors using 2D Fourier transform interpolation.
+    
+    Parameters:
+    - v0 (torch.Tensor): The first input tensor.
+    - v1 (torch.Tensor): The second input tensor.
+    - t (float): Interpolation parameter (0 <= t <= 1).
+    
+    Returns:
+    - torch.Tensor: The tensor resulting from the interpolated inverse FFT.
+    """
+    v0 = v0.to("cuda:0")
+    v1 = v1.to("cuda:0")
+
+    # Ensure the input tensors are on the same device and dtype
+    if len(v0.shape) == 1:
+        fft_v0 = torch.fft.fft(v0)
+        fft_v1 = torch.fft.fft(v1)
+        result_fft = torch.zeros_like(fft_v0)
+        
+        real_v0 = fft_v0.real
+        real_v1 = fft_v1.real
+        abs_real_v0 = real_v0.abs() 
+        abs_real_v1 = real_v1.abs() 
+
+        sign_mask = real_v0.sign() == real_v1.sign()
+        larger_values_mask = abs_real_v0 > abs_real_v1
+
+        result_fft.real[sign_mask] = (1 - t) * real_v0[sign_mask] + t * real_v1[sign_mask]
+        result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])
+
+        imag_v0 = fft_v0.imag
+        imag_v1 = fft_v1.imag
+        abs_imag_v0 = imag_v0.abs()
+        abs_imag_v1 = imag_v1.abs()
+        larger_values_mask_imag = abs_imag_v0 > abs_imag_v1
+
+        result_fft.imag[sign_mask] = (1 - t) * imag_v0[sign_mask] + t * imag_v1[sign_mask]
+        result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])
+
+        merged_tensor = torch.fft.ifft(result_fft).real  # Taking the real part
+        del v0, v1, fft_v0, fft_v1, result_fft
+        return merged_tensor
+
+    # Perform the 2D FFT on both tensors
+    fft_v0 = torch.fft.fftn(v0, dim=(-2, -1))
+    fft_v1 = torch.fft.fftn(v1, dim=(-2, -1))
+    
+    # Initialize the result FFT tensor
+    result_fft = torch.zeros_like(fft_v0)
+    
+    # Compare real parts of the coefficients
+    real_v0 = fft_v0.real
+    real_v1 = fft_v1.real
+    abs_real_v0 = real_v0.abs()
+    abs_real_v1 = real_v1.abs()
+    
+    # Create masks for where signs match and where they do not
+    sign_mask = real_v0.sign() == real_v1.sign()
+    larger_values_mask = abs_real_v0 > abs_real_v1
+    
+    # Where signs match, interpolate; where signs do not match, take the larger by magnitude
+    result_fft.real[sign_mask] = (1 - t) * real_v0[sign_mask] + t * real_v1[sign_mask]
+    result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])
+
+    del real_v0, real_v1, abs_real_v0, abs_real_v1, larger_values_mask
+    
+    # Assuming the imaginary part should be treated similarly, adjust this if not
+    imag_v0 = fft_v0.imag
+    imag_v1 = fft_v1.imag
+    abs_imag_v0 = imag_v0.abs()
+    abs_imag_v1 = imag_v1.abs()
+    larger_values_mask_imag = abs_imag_v0 > abs_imag_v1
+    
+    result_fft.imag[sign_mask] = (1 - t) * imag_v0[sign_mask] + t * imag_v1[sign_mask]
+    result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])
+
+    del imag_v0, imag_v1, abs_imag_v0, abs_imag_v1, larger_values_mask_imag, sign_mask
+    
+    # Perform the inverse FFT to go back to the spatial domain
+    merged_tensor = torch.fft.ifftn(result_fft, dim=(-2, -1)).real  # Taking the real part
+
+    del fft_v0, fft_v1, result_fft
+
+    return merged_tensor
+
+def correlate_pairs(tensors : torch.Tensor, work_device : str = "cuda:0", store_device : str = "cpu") -> torch.Tensor:
+    n = tensors.shape[0]
+    matrix = torch.zeros(n, n).to(store_device)
+    for i in range(n):
+        a = tensors[i].to(work_device)
+        for j in range(i + 1, n):
+            b = tensors[j].to(work_device)
+            matrix[i, j] = matrix[j, i] = torch.nn.functional.cosine_similarity(a, b, dim=0).nan_to_num(0).mean().item()
+            b.to(store_device)
+        a.to(store_device)
+    return matrix
+
+def least_correlated_pairs(correlation_tensor: torch.Tensor) -> Generator[Tuple[int, int, float], None, None]:
+    """
+    Generates tuples of indices and their corresponding least correlation coefficient
+    from a given correlation matrix, ensuring that once an index is used, it is no longer
+    considered in future tuples.
+
+    Args:
+        correlation_tensor (torch.Tensor): A 2D square tensor representing the correlation matrix.
+
+    Yields:
+        Tuple[int, int, float]: A tuple containing the x-index, y-index, and the correlation coefficient
+                                of the least correlated pairs in the matrix.
+    """
+    n = correlation_tensor.size(0)
+    # Create a mask to exclude diagonal and already processed elements
+    mask = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
+    
+    while torch.any(mask):
+        # Apply mask to get relevant correlations
+        valid_correlation = torch.where(mask, correlation_tensor, torch.tensor(float('inf')))
+        
+        # Find the minimum non-zero absolute correlation
+        min_val = torch.min(torch.abs(valid_correlation[valid_correlation != float('inf')]))
+        
+        # Locate the indices with the minimum correlation
+        min_indices = torch.nonzero(torch.abs(valid_correlation) == min_val, as_tuple=True)
+        if len(min_indices[0]) == 0:
+            break
+        
+        # Yield the first index pair (greedy approach) along with the correlation coefficient
+        x, y = min_indices[0][0].item(), min_indices[1][0].item()
+        coefficient = correlation_tensor[x, y].item()  # Extract the actual correlation value
+        yield (x, y, coefficient)
+        
+        # Mask out the entire row and column for both indices
+        mask[x, :] = False
+        mask[:, x] = False
+        mask[y, :] = False
+        mask[:, y] = False
+
+
+def merge_tensors_fft2_autoscale(v0: torch.Tensor, v1: torch.Tensor, t: float) -> tuple[torch.Tensor, float, float]:
+    """
+    Merges two tensors using 2D Fourier transform interpolation.
+    
+    Parameters:
+    - v0 (torch.Tensor): The first input tensor.
+    - v1 (torch.Tensor): The second input tensor.
+    - t (float): Interpolation parameter (0 <= t <= 1).
+    
+    Returns:
+    - torch.Tensor: The tensor resulting from the interpolated inverse FFT.
+    """
+    v0 = v0.to("cuda:0")
+    v1 = v1.to("cuda:0")
+
+    # Calculate norms of each tensor
+    norm_v0_t = v0.norm()
+    norm_v1_t = v1.norm()
+
+    # Scale tensors by their norms
+    v0 = v0 / norm_v0_t if norm_v0_t != 0 else v0
+    v1 = v1 / norm_v1_t if norm_v1_t != 0 else v1
+    
+    norm_v0 = norm_v0_t.item()
+    norm_v1 = norm_v1_t.item()
+    del norm_v0_t, norm_v1_t
+    
+    # Ensure the input tensors are on the same device and dtype
+    if len(v0.shape) == 1:
+        fft_v0 = torch.fft.fft(v0)
+        fft_v1 = torch.fft.fft(v1)
+        result_fft = torch.zeros_like(fft_v0)
+        
+        real_v0 = fft_v0.real
+        real_v1 = fft_v1.real
+        abs_real_v0 = real_v0.abs() 
+        abs_real_v1 = real_v1.abs() 
+
+        sign_mask = real_v0.sign() == real_v1.sign()
+        larger_values_mask = abs_real_v0 > abs_real_v1
+
+        result_fft.real[sign_mask] = (1 - t) * real_v0[sign_mask] + t * real_v1[sign_mask]
+        result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])
+
+        imag_v0 = fft_v0.imag
+        imag_v1 = fft_v1.imag
+        abs_imag_v0 = imag_v0.abs()
+        abs_imag_v1 = imag_v1.abs()
+        larger_values_mask_imag = abs_imag_v0 > abs_imag_v1
+
+        result_fft.imag[sign_mask] = (1 - t) * imag_v0[sign_mask] + t * imag_v1[sign_mask]
+        result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])
+
+        merged_tensor = torch.fft.ifft(result_fft).real  # Taking the real part
+        del v0, v1, fft_v0, fft_v1, result_fft
+        return merged_tensor, norm_v0, norm_v1
+
+    # Perform the 2D FFT on both tensors
+    fft_v0 = torch.fft.fftn(v0, dim=(-2, -1))
+    fft_v1 = torch.fft.fftn(v1, dim=(-2, -1))
+    
+    # Initialize the result FFT tensor
+    result_fft = torch.zeros_like(fft_v0)
+    
+    # Compare real parts of the coefficients
+    real_v0 = fft_v0.real
+    real_v1 = fft_v1.real
+    abs_real_v0 = real_v0.abs()
+    abs_real_v1 = real_v1.abs()
+    
+    # Create masks for where signs match and where they do not
+    sign_mask = real_v0.sign() == real_v1.sign()
+    larger_values_mask = abs_real_v0 > abs_real_v1
+    
+    # Where signs match, interpolate; where signs do not match, take the larger by magnitude
+    result_fft.real[sign_mask] = (1 - t) * real_v0[sign_mask] + t * real_v1[sign_mask]
+    result_fft.real[~sign_mask] = torch.where(larger_values_mask[~sign_mask], real_v0[~sign_mask], real_v1[~sign_mask])
+
+    del real_v0, real_v1, abs_real_v0, abs_real_v1, larger_values_mask
+    
+    # Assuming the imaginary part should be treated similarly, adjust this if not
+    imag_v0 = fft_v0.imag
+    imag_v1 = fft_v1.imag
+    abs_imag_v0 = imag_v0.abs()
+    abs_imag_v1 = imag_v1.abs()
+    larger_values_mask_imag = abs_imag_v0 > abs_imag_v1
+    
+    result_fft.imag[sign_mask] = (1 - t) * imag_v0[sign_mask] + t * imag_v1[sign_mask]
+    result_fft.imag[~sign_mask] = torch.where(larger_values_mask_imag[~sign_mask], imag_v0[~sign_mask], imag_v1[~sign_mask])
+
+    del imag_v0, imag_v1, abs_imag_v0, abs_imag_v1, larger_values_mask_imag, sign_mask
+    
+    # Perform the inverse FFT to go back to the spatial domain
+    merged_tensor = torch.fft.ifftn(result_fft, dim=(-2, -1)).real  # Taking the real part
+
+    del fft_v0, fft_v1, result_fft
+
+    return merged_tensor, norm_v0, norm_v1

ztrain/util.py

@@ -0,0 +1,37 @@
+# ztrain/util.py
+# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted
+
+import contextlib
+import torch
+
+
+@contextlib.contextmanager
+def cuda_memory_profiler(display : str = True):
+    """
+    A context manager for profiling CUDA memory usage in PyTorch.
+    """
+    if display is False:
+        yield
+        return
+    
+    if not torch.cuda.is_available():
+        print("CUDA is not available, skipping memory profiling")
+        yield
+        return
+    
+    torch.cuda.reset_peak_memory_stats()
+    torch.cuda.synchronize()
+    start_memory = torch.cuda.memory_allocated()
+
+    try:
+        yield
+    finally:
+        torch.cuda.synchronize()
+        end_memory = torch.cuda.memory_allocated()
+        print(f"Peak memory usage: {torch.cuda.max_memory_allocated() / (1024 ** 2):.2f} MB")
+        print(f"Memory allocated at start: {start_memory / (1024 ** 2):.2f} MB")
+        print(f"Memory allocated at end: {end_memory / (1024 ** 2):.2f} MB")
+        print(f"Net memory change: {(end_memory - start_memory) / (1024 ** 2):.2f} MB")
+
+def get_device():
+    return torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")