model.model.model.model.model

mosaicml · Apr 22, 2024 · be5105d · be5105d
1 parent e1e59fe
commit be5105d
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 11 deletions.
diff --git a/llmfoundry/models/hf/hf_causal_lm.py b/llmfoundry/models/hf/hf_causal_lm.py
@@ -246,9 +246,9 @@ def _autoset_attn_implementation_monkeypatch(
         if model.config.tie_word_embeddings and resolved_init_device == 'meta':
             model.tie_weights()
 
-        peft_config = None
+        peft_config_object = None
         if peft_config is not None:
-            peft_config = self._get_peft_config(peft_config)
+            peft_config_object = self._get_peft_config(peft_config)
 
         if pretrained_lora_id_or_path is not None:
             if not peft_installed:
@@ -266,7 +266,7 @@ def _autoset_attn_implementation_monkeypatch(
             metrics=train_metrics,
             eval_metrics=eval_metrics,
             init_device=init_device,
-            peft_config=peft_config,
+            peft_config=peft_config_object,
         )
 
     @staticmethod

diff --git a/tests/models/hf/test_fsdp_weight_tying.py b/tests/models/hf/test_fsdp_weight_tying.py
@@ -7,10 +7,8 @@
 import pytest
 from composer import Trainer
 from composer.models.huggingface import maybe_get_underlying_model
-from omegaconf import OmegaConf as om
 
 from llmfoundry.utils.builders import build_composer_model, build_tokenizer
-from llmfoundry.utils.config_utils import to_dict_recursive
 
 
 @pytest.mark.world_size(2)
@@ -48,7 +46,6 @@ def test_fsdp_weight_tying(peft_config: Optional[dict], tmp_path: pathlib.Path,
 
     assert model_cfg is not None
     assert tokenizer_name is not None
-    model_cfg = om.create(model_cfg)
     if peft_config is not None:
         model_cfg['peft_config'] = peft_config
 
@@ -71,12 +68,12 @@ def test_fsdp_weight_tying(peft_config: Optional[dict], tmp_path: pathlib.Path,
     name = model_cfg.pop('name')
     original_model = build_composer_model(
         name=name,
-        cfg=to_dict_recursive(model_cfg),
+        cfg=model_cfg,
         tokenizer=tokenizer,
     )
 
     underlying_model = maybe_get_underlying_model(original_model.model)
-    lm_head = underlying_model.lm_head if peft_config is None else underlying_model.lm_head
+    lm_head = underlying_model.lm_head
     embedding_layer = underlying_model.model.embed_tokens if peft_config is None else underlying_model.model.embed_tokens
 
     lm_head_id = id(lm_head.weight)
@@ -93,8 +90,8 @@ def test_fsdp_weight_tying(peft_config: Optional[dict], tmp_path: pathlib.Path,
     )
 
     model = trainer.state.model
-    lm_head = model.lm_head if peft_config is None else model.base_model.lm_head
-    embedding_layer = model.embed_tokens if peft_config is None else model.base_model.model.embed_tokens
+    lm_head = model.model.lm_head if peft_config is None else model.model.base_model.model.lm_head
+    embedding_layer = model.model.model.embed_tokens if peft_config is None else model.model.base_model.model.model.embed_tokens
 
     lm_head_id = id(lm_head.weight)
     embedding_layer_id = id(embedding_layer.weight)

diff --git a/tests/models/hf/test_hf_peft_wrapping.py b/tests/models/hf/test_hf_peft_wrapping.py
@@ -99,7 +99,8 @@ def test_lora_mixed_init(peft_config: Optional[dict], tmp_path: pathlib.Path,
     )
 
     model = trainer.state.model
-    underlying_model = model.base_model
+    underlying_model = model.model.base_model.model.model
+    # assert False, f"underlying_model: {underlying_model}"
     lora_A = underlying_model.layers[0].self_attn.q_proj.lora_A['default']
     lora_B = underlying_model.layers[0].self_attn.q_proj.lora_B['default']