Upload 2 files

config.json

@@ -0,0 +1,16 @@
+{
+  "architectures": [
+    "TestLakeForConditionalGeneration"
+  ],
+  "model_type": "testlake",
+  "torch_dtype": "torch.bfloat16",
+  "hidden_size": 1344,
+  "num_hidden_layers": 3,
+  "num_attention_heads": 4,
+  "intermediate_size": 1344,
+  "head_dim": 336,
+  "vocab_size": 65552,
+  "image_feature_dim": 1344,
+  "num_patches": 196,
+  "vision_num_layers": 6
+}

main.py

@@ -0,0 +1,219 @@
+import torch
+from safetensors.torch import save_file, load_file
+import sys
+import json
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def create_bitnet_safetensors(output_file="model_1.safetensors", 
+                             dtype=torch.bfloat16, 
+                             num_heads=4, 
+                             num_layers=4, 
+                             image_feature_dim=256, 
+                             device=device,
+                             quantize=True):
+    """
+    Creates a safetensors file with initialized BitNet model weights.
+    """
+    tensors = {}
+
+    # These values reduce model size; adjust as needed
+    embedding_dim = 5376 // 4
+    intermediate_size = 21504 // 4
+    head_dim = embedding_dim // num_heads
+
+    # Text Encoder Components
+    # Text embedding
+    tensors["language_model.model.embed_tokens.weight"] = torch.randn(262208 // 4, embedding_dim, dtype=dtype, device=device) * 0.02
+
+    # Text layers
+    for layer_idx in range(num_layers):
+        layer_prefix = f"language_model.model.layers.{layer_idx}"
+
+        # Layer normalization
+        tensors[f"{layer_prefix}.pre_attn_layernorm.weight"] = torch.ones(embedding_dim, dtype=dtype, device=device)
+        tensors[f"{layer_prefix}.post_mlp_layernorm.weight"] = torch.ones(embedding_dim, dtype=dtype, device=device)
+
+        # Self-attention projections (Q, K, V)
+        tensors[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
+
+        # Self-attention output projection
+        tensors[f"{layer_prefix}.self_attn.out_proj.weight"] = torch.randn(embedding_dim, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
+
+        # MLP sub-block: first linear layer
+        tensors[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(intermediate_size, embedding_dim, dtype=dtype, device=device) * (embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(intermediate_size, dtype=dtype, device=device)
+
+        # MLP sub-block: second linear layer
+        tensors[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(embedding_dim, intermediate_size, dtype=dtype, device=device) * (intermediate_size ** -0.5)
+        tensors[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(embedding_dim, dtype=dtype, device=device)
+
+    # Vision Encoder Components
+    patch_size = 16
+    image_size = 224
+    num_patches = (image_size // patch_size) ** 2
+    vision_embedding_dim = embedding_dim
+
+    # Patch Embedding
+    tensors["vision_encoder.patch_embedding.weight"] = torch.randn(vision_embedding_dim, 3, patch_size, patch_size, dtype=dtype, device=device) * 0.02
+    tensors["vision_encoder.patch_embedding.bias"] = torch.zeros(vision_embedding_dim, dtype=dtype, device=device)
+
+    # Positional Embeddings
+    tensors["vision_encoder.position_embeddings"] = torch.randn(1, num_patches + 1, vision_embedding_dim, dtype=dtype, device=device) * 0.02
+
+    # Class Token
+    tensors["vision_encoder.cls_token"] = torch.randn(1, 1, vision_embedding_dim, dtype=dtype, device=device) * 0.02
+
+    # Vision Transformer Layers
+    for layer_idx in range(num_layers):
+        layer_prefix = f"vision_encoder.transformer.layers.{layer_idx}"
+
+        # Pre-Attention Layer Normalization
+        tensors[f"{layer_prefix}.pre_attn_layernorm.weight"] = torch.ones(vision_embedding_dim, dtype=dtype, device=device)
+
+        # Self-Attention projections (Q, K, V)
+        tensors[f"{layer_prefix}.self_attn.q_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.k_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.self_attn.v_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
+
+        # Self-Attention output projection
+        tensors[f"{layer_prefix}.self_attn.out_proj.weight"] = torch.randn(vision_embedding_dim, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
+
+        # MLP sub-block: first linear layer
+        tensors[f"{layer_prefix}.mlp.fc1.weight"] = torch.randn(vision_embedding_dim, intermediate_size, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.mlp.fc1.bias"] = torch.zeros(intermediate_size, dtype=dtype, device=device)
+
+        # MLP sub-block: second linear layer
+        tensors[f"{layer_prefix}.mlp.fc2.weight"] = torch.randn(intermediate_size, vision_embedding_dim, dtype=dtype, device=device) * (vision_embedding_dim ** -0.5)
+        tensors[f"{layer_prefix}.mlp.fc2.bias"] = torch.zeros(vision_embedding_dim, dtype=dtype, device=device)
+
+        # Post-MLP Layer Normalization
+        tensors[f"{layer_prefix}.post_mlp_layernorm.weight"] = torch.ones(vision_embedding_dim, dtype=dtype, device=device)
+
+    total_params = sum(tensor.numel() * tensor.element_size() for tensor in tensors.values())
+    total_params_gb = total_params / (1024 ** 3)
+
+    print(f"Estimated total parameter size: {total_params_gb:.2f} GB ({dtype}).")
+
+    proceed = input("Do you want to proceed and create the safetensors file? (y/n): ").lower()
+
+    if proceed == 'y':
+        if quantize:
+            # Apply 1.58-bit quantization
+            quantized_tensors = {}
+            for key, tensor in tensors.items():
+                # Normalize the tensor
+                max_val = torch.max(torch.abs(tensor))
+                normalized = tensor / max_val
+                
+                # Quantize to {-1, 0, 1}
+                quantized = torch.where(normalized > 0.5, torch.ones_like(normalized), 
+                                      torch.where(normalized < -0.5, -torch.ones_like(normalized), torch.zeros_like(normalized)))
+                
+                quantized_tensors[key] = quantized
+                print(f"Quantized {key} to 1.58-bit precision.")
+            
+            tensors = quantized_tensors
+
+        tensors = {k: v.cpu() for k, v in tensors.items()}
+        save_file(tensors, output_file)
+        print(f"BitNet safetensors file created: {output_file}")
+    else:
+        print("Safetensors file creation cancelled.")
+        sys.exit()
+
+def quantize_model(model_file, output_file, device=device):
+    """
+    Quantizes the given model to 1.58-bit precision.
+    """
+    try:
+        # Load the model
+        tensors = load_file(model_file, device=device)
+        
+        # Quantize each tensor
+        quantized_tensors = {}
+        for key, tensor in tensors.items():
+            # Normalize the tensor
+            max_val = torch.max(torch.abs(tensor))
+            normalized = tensor / max_val
+            
+            # Quantize to {-1, 0, 1}
+            quantized = torch.where(normalized > 0.5, torch.ones_like(normalized), 
+                                  torch.where(normalized < -0.5, -torch.ones_like(normalized), torch.zeros_like(normalized)))
+            
+            quantized_tensors[key] = quantized
+            print(f"Quantized {key} to 1.58-bit precision.")
+        
+        # Save the quantized model
+        save_file(quantized_tensors, output_file)
+        print(f"Quantized model saved to {output_file}")
+        
+    except FileNotFoundError:
+        print(f"Error: Model file '{model_file}' not found.")
+    except Exception as e:
+        print(f"An error occurred during quantization: {e}")
+
+def analyze_model_file(model_file):
+    """
+    Analyzes a safetensors model file and prints/saves the configuration.
+    """
+    try:
+        tensors = load_file(model_file, device="cpu")
+        config = {}
+
+        # Basic model architecture details
+        config["architectures"] = ["TestLakeForConditionalGeneration"]
+        config["model_type"] = "testlake"
+        config["torch_dtype"] = str(tensors["language_model.model.embed_tokens.weight"].dtype)
+
+        # Extract dimensions for language model
+        config["hidden_size"] = tensors["language_model.model.embed_tokens.weight"].shape[1]
+        config["num_hidden_layers"] = len([key for key in tensors if "language_model.model.layers." in key]) // 12
+        num_heads = tensors["language_model.model.layers.0.self_attn.q_proj.weight"].shape[0] // (config["hidden_size"] // 4)
+        config["num_attention_heads"] = num_heads
+        config["intermediate_size"] = tensors["language_model.model.layers.0.mlp.fc1.weight"].shape[1]
+        config["head_dim"] = config["hidden_size"] // config["num_attention_heads"]
+
+        # Extract vocabulary and vision details
+        config["vocab_size"] = tensors["language_model.model.embed_tokens.weight"].shape[0]
+        config["image_feature_dim"] = tensors["vision_encoder.patch_embedding.weight"].shape[0]
+        config["num_patches"] = ((224 // 16) ** 2)
+        config["vision_num_layers"] = len([key for key in tensors if "vision_encoder.transformer.layers." in key]) // 6
+
+        print(json.dumps(config, indent=2))
+
+        # Save to config.json
+        with open("config.json", "w") as f:
+            json.dump(config, f, indent=2)
+
+        print("Model analysis completed and saved to config.json")
+
+    except FileNotFoundError:
+        print(f"Error: Model file '{model_file}' not found.")
+    except Exception as e:
+        print(f"An error occurred during model analysis: {e}")
+
+if __name__ == "__main__":
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    # Set the desired data type for training (BFloat16)
+    dtype_input = input("Enter the desired data type (bfloat16, float16, float32, etc.): ").lower()
+    if dtype_input == "bfloat16":
+        dtype = torch.bfloat16
+    elif dtype_input == "float16":
+        dtype = torch.float16
+    elif dtype_input == "float32":
+        dtype = torch.float32
+    elif dtype_input == "float8":
+        dtype = torch.float8
+    else:
+        print("Invalid data type. Using bfloat16 as default.")
+        dtype = torch.bfloat16
+
+    # Create the initial model with BFloat16 precision and quantize it
+    create_bitnet_safetensors(dtype=dtype, device=device, quantize=True)
+
+    # Analyze the created model file
+    analyze_model_file("model_1.safetensors")