updt pr comment

llmfoundry/models/mpt/modeling_mpt.py

@@ -580,7 +580,7 @@ def __init__(self, config: MPTConfig):
         self.transformer: MPTModel = MPTModel(config)
 
         self.lm_head = None
-        if config.tie_word_embeddings is False:
+        if not config.tie_word_embeddings:
             self.lm_head = nn.Linear(
                 config.d_model,
                 config.vocab_size,
@@ -882,7 +882,7 @@ def flops_per_batch(self, batch: Mapping) -> int:
 
         bs, msl = batch['input_ids'].shape[0:2]
         params = self.n_active_params
-        if self.model.transformer.config.tie_word_embeddings is False:
+        if not self.model.transformer.config.tie_word_embeddings:
             # embedding layers are lookup tables, therefore are not counted in the FLOP computation
             params -= self.model.transformer.wte.weight.numel()
         params_flops_per_token = 2 * params

tests/test_model.py

@@ -954,9 +954,7 @@ def check_hf_model_equivalence(model1: PreTrainedModel,
         torch.testing.assert_close(p1, p2)
 
 
-@pytest.mark.parametrize('tie_word_embeddings', [True, False])
-def test_save_from_pretrained(tie_word_embeddings: bool,
-                              tmp_path: pathlib.Path):
+def test_save_from_pretrained(tmp_path: pathlib.Path):
     # Test that MPT can be used with the HuggingFace
     # save_pretrained/from_pretrained api.
     hf_config = MPTConfig(
@@ -971,12 +969,10 @@ def test_save_from_pretrained(tie_word_embeddings: bool,
         attn_config={
             'attn_impl': 'torch',
         },
-        tie_word_embeddings=tie_word_embeddings,
     )
     mpt = MPTForCausalLM(hf_config)
 
     mpt.save_pretrained(tmp_path / 'test-save-pretrained')
-    print(tmp_path / 'test-save-pretrained')
     mpt2 = MPTForCausalLM.from_pretrained(tmp_path / 'test-save-pretrained')
 
     check_hf_model_equivalence(mpt, mpt2)