diff --git a/README.md b/README.md index a126d92..fc571bd 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,12 @@ Currently `flow-merge` supports most of the popular and proven merge methods. | `qwen` | `QwenForCausalLM` | | `mistral` | `MistralForCausalLM` | | `llama` | `LlamaForCausalLM` | +| `gemma` | `GemmaForCausalLM` | +| `phi` | `PhiForCausalLM` | +| `phi3` | `Phi3ForCausalLM` | + + + > 📢 _We plan to support many models and architectures more, including encoder models such as BERT-Family models too._ diff --git a/flow_merge/data/architectures/gemma.json b/flow_merge/data/architectures/gemma.json new file mode 100644 index 0000000..22fa243 --- /dev/null +++ b/flow_merge/data/architectures/gemma.json @@ -0,0 +1,72 @@ +{ + "model_type": "gemma", + "architectures": [ + "GemmaForCausalLM" + ], + "num_layers_config_key": "num_hidden_layers", + "weights": [ + { + "name": "model.embed_tokens.weight", + "type": "embed_tokens", + "layer_type": "embedding" + }, + { + "name": "model.layers.{layer_index}.input_layernorm.weight", + "type": "input_layernorm", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.q_proj.weight", + "type": "self_attn", + "projection": "q_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.k_proj.weight", + "type": "self_attn", + "projection": "k_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.v_proj.weight", + "type": "self_attn", + "projection": "v_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.o_proj.weight", + "type": "self_attn", + "projection": "o_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.post_attention_layernorm.weight", + "type": "post_attention_layernorm", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.gate_proj.weight", + "type": "mlp", + "projection": "gate_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.up_proj.weight", + "type": "mlp", + "projection": "up_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.down_proj.weight", + "type": "mlp", + "projection": "down_proj", + "layer_type": "decoder" + }, + { + "name": "model.norm.weight", + "type": "norm", + "layer_type": "post_norm" + } + ], + "residual": {} +} \ No newline at end of file diff --git a/flow_merge/data/architectures/phi.json b/flow_merge/data/architectures/phi.json new file mode 100644 index 0000000..bb4b76b --- /dev/null +++ b/flow_merge/data/architectures/phi.json @@ -0,0 +1,117 @@ +{ + "model_type": "phi", + "architectures": [ + "PhiForCausalLM" + ], + "num_layers_config_key": "num_hidden_layers", + "weights": [ + { + "name": "model.embed_tokens.weight", + "type": "embed_tokens", + "layer_type": "embedding" + }, + { + "name": "model.layers.{layer_index}.input_layernorm.weight", + "type": "input_layernorm", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.input_layernorm.bias", + "type": "input_layernorm", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.q_proj.weight", + "type": "self_attn", + "projection": "q_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.q_proj.bias", + "type": "self_attn", + "projection": "q_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.k_proj.weight", + "type": "self_attn", + "projection": "k_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.k_proj.bias", + "type": "self_attn", + "projection": "k_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.v_proj.weight", + "type": "self_attn", + "projection": "v_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.v_proj.bias", + "type": "self_attn", + "projection": "v_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.dense.weight", + "type": "self_attn", + "projection": "dense", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.dense.bias", + "type": "self_attn", + "projection": "dense", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.fc1.weight", + "type": "mlp", + "projection": "fc1", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.fc1.bias", + "type": "mlp", + "projection": "fc1", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.fc2.weight", + "type": "mlp", + "projection": "fc2", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.fc2.bias", + "type": "mlp", + "projection": "fc2", + "layer_type": "decoder" + }, + { + "name": "model.final_layernorm.weight", + "type": "final_layernorm", + "layer_type": "layernorm" + }, + { + "name": "model.final_layernorm.bias", + "type": "final_layernorm", + "layer_type": "layernorm" + }, + { + "name": "lm_head.weight", + "type": "lm_head", + "layer_type": "head" + }, + { + "name": "lm_head.bias", + "type": "lm_head", + "layer_type": "head" + } + ], + "residual": {} +} \ No newline at end of file diff --git a/flow_merge/data/architectures/phi3.json b/flow_merge/data/architectures/phi3.json new file mode 100644 index 0000000..d8f9891 --- /dev/null +++ b/flow_merge/data/architectures/phi3.json @@ -0,0 +1,59 @@ +{ + "model_type": "phi3", + "architectures": [ + "Phi3ForCausalLM" + ], + "num_layers_config_key": "num_hidden_layers", + "weights": [ + { + "name": "model.embed_tokens.weight", + "type": "embed_tokens", + "layer_type": "embedding" + }, + { + "name": "model.layers.{layer_index}.input_layernorm.weight", + "type": "input_layernorm", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.o_proj.weight", + "type": "self_attn", + "projection": "o_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.self_attn.qkv_proj.weight", + "type": "self_attn", + "projection": "qkv_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.gate_up_proj.weight", + "type": "mlp", + "projection": "gate_up_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.mlp.down_proj.weight", + "type": "mlp", + "projection": "down_proj", + "layer_type": "decoder" + }, + { + "name": "model.layers.{layer_index}.post_attention_layernorm.weight", + "type": "post_attention_layernorm", + "layer_type": "decoder" + }, + { + "name": "model.norm.weight", + "type": "norm", + "layer_type": "post_norm" + }, + { + "name": "lm_head.weight", + "type": "lm_head", + "layer_type": "head" + } + ], + "residual": {} +} diff --git a/flow_merge/lib/architecture.py b/flow_merge/lib/architecture.py index cd9cfc9..a4b516b 100644 --- a/flow_merge/lib/architecture.py +++ b/flow_merge/lib/architecture.py @@ -10,9 +10,9 @@ import flow_merge.data.architectures ArchitectureType = Enum( - "ArchitectureType", ["MistralForCausalLM", "LlamaForCausalLM", "Qwen2ForCausalLM"] + "ArchitectureType", ["MistralForCausalLM", "LlamaForCausalLM", "Qwen2ForCausalLM", "PhiForCausalLM", "Phi3ForCausalLM"] ) -ModelType = Enum("ModelType", ["mistral", "llama", "qwen-1.5"]) +ModelType = Enum("ModelType", ["mistral", "llama", "qwen-1.5", "phi", "phi3"]) ModelWeightType = Enum( "ModelWeightType", [ @@ -22,15 +22,16 @@ "embed_tokens", "norm", "lm_head", - "post_attention_layernorm", + "post_attention_layernorm", + "final_layernorm", ], ) ProjectionType = Enum( "ProjectionType", - ["v_proj", "q_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], + ["v_proj", "q_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "dense", "fc1", "fc2", "qkv_proj", "gate_up_proj"], ) ModelWeightLayerType = Enum( - "ModelWeightLayerType", ["decoder", "embedding", "head", "post_norm"] + "ModelWeightLayerType", ["decoder", "embedding", "head", "post_norm", "layernorm"] ) diff --git a/flow_merge/lib/merge_config.py b/flow_merge/lib/merge_config.py index e2dbe1f..7ba9803 100644 --- a/flow_merge/lib/merge_config.py +++ b/flow_merge/lib/merge_config.py @@ -132,8 +132,8 @@ def __init__( self.directory_settings: DirectorySettings = data.directory_settings self.hf_hub_settings: HfHubSettings = data.hf_hub_settings self.device = self.select_device() - self.models: List[Model] = self.create_models() - self.base_model: Model = self.create_base_model() + self.models: List[Model] = self.create_models(trust_remote_code=self.hf_hub_settings.trust_remote_code) + self.base_model: Model = self.create_base_model(trust_remote_code=self.hf_hub_settings.trust_remote_code) self.method_config = self._extract_and_set_weights() @staticmethod @@ -171,18 +171,18 @@ def select_device(self) -> str: return "cpu" return device - def create_models(self) -> List[Model]: + def create_models(self, trust_remote_code: bool) -> List[Model]: # observe that models list doesn't contain the base_model by design # access base_model: Model by accessing base_model variable models_data = [ m for m in self.data.models if m.path_or_id != self.data.base_model ] - return [Model.from_path(model_data.path_or_id) for model_data in models_data] + return [Model.from_path(model_data.path_or_id, trust_remote_code=trust_remote_code) for model_data in models_data] # TODO - this should take trust_remote_code - def create_base_model(self) -> Model: + def create_base_model(self, trust_remote_code: bool) -> Model: if not self.data.base_model: first_model_from_list = self.data.models[0] - return Model.from_path(first_model_from_list.path_or_id) + return Model.from_path(first_model_from_list.path_or_id, trust_remote_code=trust_remote_code) # TODO - this should take trust_remote_code else: base_model = next( ( @@ -197,7 +197,7 @@ def create_base_model(self) -> Model: f"Base model '{self.data.base_model}' not found in the list of " + f"models {[model.path_or_id for model in self.data.models]}." ) - return Model.from_path(base_model.path_or_id) + return Model.from_path(base_model.path_or_id, trust_remote_code=trust_remote_code) def _extract_and_set_weights( self,