Upload 32to31.py

swapper/32to31.py

@@ -0,0 +1,182 @@
+import json
+import os
+import re
+from huggingface_hub import snapshot_download
+
+import torch
+from safetensors import safe_open
+from transformers import AutoProcessor, MllamaForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM, AutoConfig
+
+#total_layers=80  # 70B model has 80 layers
+total_layers=32   #  8B model has 32 layers
+
+#cross_attention_layers = [3, 8, 13, 18, 23, 28, 33, 38, 43, 48, 53, 58, 63, 68, 73, 78, 83, 88, 93, 98] # 90B
+cross_attention_layers = [3, 8, 13, 18, 23, 28, 33, 38] # 11b
+
+# Update paths - switch source and target
+target_model = "meta-llama/Llama-3.1-8B-Instruct"
+print(f"Target model: {target_model}")
+
+source_model = "meta-llama/Llama-3.2-11B-Vision-Instruct"
+print(f"Source model: {source_model}")
+
+def create_inverse_layer_mapping(total_layers=total_layers, cross_attn_layers=cross_attention_layers):
+    """
+    Creates a mapping from 90B/11B layer indices to 70B/8B layer indices.
+    """
+    mapping = {}
+    removed_layers = []
+
+    #for i in range(100):  # 90B has 100 layers (80 + 20 cross-attention layers)
+    for i in range(40):  # 11B has 40 layers (32 + 8 cross-attention layers)
+        if i not in cross_attn_layers and len(mapping) < total_layers:
+            mapping[i] = len(mapping)
+        else:
+            removed_layers.append(i)
+    return mapping, removed_layers
+
+def load_sharded_state_dict(model_id):
+    """
+    Load a sharded state dict from either a local directory or a Hugging Face model ID.
+    
+    Args:
+        model_id: Either a local path or a Hugging Face model ID (e.g., "meta-llama/Llama-2-7b")
+    
+    Returns:
+        dict: The loaded state dictionary
+    """
+    # Check if model_id is a local path
+    if os.path.isdir(model_id):
+        model_dir = model_id
+    else:
+        # If not local, assume it's a Hugging Face model ID and download it
+        print(f"Downloading model from Hugging Face: {model_id}")
+        model_dir = snapshot_download(
+            model_id,
+            allow_patterns=["*.safetensors*", "*.json"],
+            ignore_patterns=["*.bin", "*.md", "*.py"]
+        )
+    
+    # Load the index file
+    index_file = os.path.join(model_dir, 'model.safetensors.index.json')
+    if not os.path.exists(index_file):
+        raise FileNotFoundError(f"Could not find index file: {index_file}")
+    
+    with open(index_file, 'r') as f:
+        index_data = json.load(f)
+    
+    weight_map = index_data['weight_map']
+    state_dict = {}
+    shard_to_params = {}
+    
+    # Group parameters by shard file
+    for param_name, shard_file in weight_map.items():
+        if shard_file not in shard_to_params:
+            shard_to_params[shard_file] = []
+        shard_to_params[shard_file].append(param_name)
+    
+    # Load parameters from each shard
+    for shard_file, params_in_shard in shard_to_params.items():
+        shard_path = os.path.join(model_dir, shard_file)
+        with safe_open(shard_path, framework="pt", device="cpu") as f:
+            for name in params_in_shard:
+                state_dict[name] = f.get_tensor(name)
+    
+    return state_dict
+
+def compare_model_states(model, new_state_dict):
+    current_state = model.state_dict()
+    unchanged_params = []
+    changed_params = []
+    missing_params = []
+
+    for name, param in current_state.items():
+        if name not in new_state_dict:
+            missing_params.append(name)
+        elif torch.equal(param, new_state_dict[name]):
+            unchanged_params.append(name)
+        else:
+            sum_abs_diff = torch.sum(torch.abs(param - new_state_dict[name]))
+            changed_params.append({'name': name, 'sum_abs_diff': sum_abs_diff.item()})
+
+    return {
+        'unchanged': unchanged_params,
+        'changed': changed_params,
+        'missing': missing_params
+    }
+
+
+layer_mapping, removed_layers = create_inverse_layer_mapping()
+
+# Load source (90B) state dict
+source_state_dict = load_sharded_state_dict(source_model)
+
+# Create new state dict for target model (70B)
+target_state_dict = {}
+
+# Convert parameter names and copy tensors
+for name, param in source_state_dict.items():
+    # Skip parameters that aren't part of the language model layers
+    if not (name.startswith('language_model.model.layers.') or 
+            name == 'language_model.model.embed_tokens.weight' or
+            name == 'language_model.lm_head.weight' or
+            name == 'language_model.model.norm.weight'):
+        continue
+    
+    if name.startswith('language_model.model.layers.'):
+        # Handle layer parameters
+        layer_match = re.match(r'language_model\.model\.layers\.(\d+)\.(.+)', name)
+        if layer_match:
+            source_layer = int(layer_match.group(1))
+            if source_layer in layer_mapping:
+                target_layer = layer_mapping[source_layer]
+                new_name = f'model.layers.{target_layer}.{layer_match.group(2)}'
+                target_state_dict[new_name] = param
+    elif name == 'language_model.lm_head.weight':
+        # Handle lm_head weight
+        target_state_dict['lm_head.weight'] = param
+    elif name == 'language_model.model.embed_tokens.weight':
+        # Handle embeddings - keep original vocab size for 70B model
+        original_embed_size = 128256
+        target_state_dict['model.embed_tokens.weight'] = param[:original_embed_size, :]
+    elif name == 'language_model.model.norm.weight':
+        # Handle model norm weight
+        target_state_dict['model.norm.weight'] = param
+
+
+#write target_state_dict keys to file for verification
+with open('target_state_dict.txt', 'w') as f:
+    f.write('\n'.join(target_state_dict.keys()))
+
+config = AutoConfig.from_pretrained(target_model)
+
+model = AutoModelForCausalLM.from_pretrained(
+    None,
+    config=config,
+    state_dict = target_state_dict,
+    torch_dtype=torch.bfloat16,
+)
+
+'''
+
+origmodel = AutoModelForCausalLM.from_pretrained(
+    target_model,
+    torch_dtype=torch.bfloat16,
+)
+
+result = compare_model_states(model, origmodel.state_dict())
+print("Unchanged parameters:", len(result['unchanged']))
+print("Changed parameters:", len(result['changed']))
+print("Missing parameters:", len(result['missing']))
+
+#write result to file
+with open('result.txt', 'w') as f:
+    f.write(json.dumps(result, indent=2))
+
+'''
+
+processor = AutoTokenizer.from_pretrained(target_model) #8b/70b
+#processor = AutoProcessor.from_pretrained(source_model) #11b/90b
+
+model.save_pretrained("Llama-3.2-8B-extracted")
+processor.save_pretrained("Llama-3.2-8B-extracted")