Combine tests for phi-1.5 and phi-2 into one

Lightning-AI · Apr 23, 2024 · de45214 · de45214
1 parent 8a8b698
commit de45214
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 89 deletions.
diff --git a/tests/test_convert_lit_checkpoint.py b/tests/test_convert_lit_checkpoint.py
@@ -222,49 +222,13 @@ def test_against_original_open_llama_3b():
 
 
 @torch.inference_mode()
-def test_against_hf_phi_1_5():
+@pytest.mark.parametrize("model_name", ("phi-1_5", "phi-2"))
+def test_against_hf_phi(model_name):
     from transformers.models.phi.configuration_phi import PhiConfig
     from transformers.models.phi.modeling_phi import PhiForCausalLM
 
     ours_config = Config.from_name(
-        "phi-1_5", padded_vocab_size=10000, n_layer=2, n_head=4, n_embd=256, rotary_percentage=0.5
-    )
-    T = 5
-    theirs_config = PhiConfig(
-        vocab_size=ours_config.padded_vocab_size,
-        max_position_embeddings=ours_config.block_size,
-        hidden_size=ours_config.n_embd,
-        intermediate_size=ours_config.intermediate_size,
-        num_attention_heads=ours_config.n_head,
-        num_hidden_layers=ours_config.n_layer,
-        partial_rotary_factor=ours_config.rotary_percentage,
-    )
-
-    ours_model = GPT(ours_config)
-    ours_state_dict = ours_model.state_dict()
-    theirs_state_dict = {}
-    copy_weights_phi(ours_config, theirs_state_dict, ours_state_dict)
-    theirs_model = PhiForCausalLM(theirs_config)
-    # strict=False because we don't save the rotary embeddings inv frequency
-    keys = theirs_model.load_state_dict(theirs_state_dict, strict=False)
-    assert not keys.unexpected_keys
-    assert all("inv_freq" in k for k in keys.missing_keys)
-
-    # test end to end
-    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32)
-    assert x.size(1) == T
-    ours_y = ours_model(x)
-    theirs_y = theirs_model(x)["logits"]
-    torch.testing.assert_close(ours_y, theirs_y)
-
-
-@torch.inference_mode()
-def test_against_hf_phi_2():
-    from transformers.models.phi.configuration_phi import PhiConfig
-    from transformers.models.phi.modeling_phi import PhiForCausalLM
-
-    ours_config = Config.from_name(
-        "phi-2", padded_vocab_size=10000, n_layer=2, n_head=4, n_embd=256, rotary_percentage=0.5
+        model_name, padded_vocab_size=10000, n_layer=2, n_head=4, n_embd=256, rotary_percentage=0.5
     )
     T = 5
     theirs_config = PhiConfig(

diff --git a/tests/test_model.py b/tests/test_model.py
@@ -267,6 +267,7 @@ def test_against_hf_llama_2_and_3(ours_kwargs, device, dtype):
 
 
 @torch.inference_mode()
+@pytest.mark.parametrize("model_name", ("phi-1_5", "phi-2"))
 @pytest.mark.parametrize(
     ("device", "dtype"),
     [
@@ -278,62 +279,14 @@ def test_against_hf_llama_2_and_3(ours_kwargs, device, dtype):
         ),
     ],
 )
-def test_against_hf_phi_1_5(device, dtype):
+def test_against_hf_phi(model_name, device, dtype):
     from transformers.models.phi.configuration_phi import PhiConfig
     from transformers.models.phi.modeling_phi import PhiForCausalLM
 
     torch.set_default_dtype(dtype)
 
     ours_config = Config.from_name(
-        "phi-1_5", padded_vocab_size=10000, n_layer=2, n_head=4, n_embd=256, rotary_percentage=0.5
-    )
-    T = 5
-    theirs_config = PhiConfig(
-        vocab_size=ours_config.padded_vocab_size,
-        max_position_embeddings=ours_config.block_size,
-        hidden_size=ours_config.n_embd,
-        intermediate_size=ours_config.intermediate_size,
-        num_attention_heads=ours_config.n_head,
-        num_hidden_layers=ours_config.n_layer,
-        partial_rotary_factor=ours_config.rotary_percentage,
-        torch_dtype=dtype,
-    )
-
-    theirs_model = PhiForCausalLM(theirs_config).to(device)
-    theirs_state_dict = theirs_model.state_dict()
-    state_dict = {}
-    copy_weights_phi(ours_config, {}, state_dict, theirs_state_dict)
-    ours_model = GPT(ours_config).to(device)
-    ours_model.load_state_dict(state_dict)
-
-    # test end to end
-    x = torch.tensor([[9856, 23, 491, 1536, 304]], dtype=torch.int32, device=device)
-    assert x.size(1) == T
-    ours_y = ours_model(x)
-    theirs_y = theirs_model(x)["logits"].to(dtype)  # HF converts logits to float
-    torch.testing.assert_close(ours_y, theirs_y)
-
-
-@torch.inference_mode()
-@pytest.mark.parametrize(
-    ("device", "dtype"),
-    [
-        (torch.device("cpu"), torch.float32),
-        pytest.param(
-            torch.device("cuda"),
-            torch.float16,
-            marks=[pytest.mark.xfail(raises=AssertionError, strict=False), RunIf(min_cuda_gpus=1)],
-        ),
-    ],
-)
-def test_against_hf_phi_2(device, dtype):
-    from transformers.models.phi.configuration_phi import PhiConfig
-    from transformers.models.phi.modeling_phi import PhiForCausalLM
-
-    torch.set_default_dtype(dtype)
-
-    ours_config = Config.from_name(
-        "phi-2", padded_vocab_size=10000, n_layer=2, n_head=4, n_embd=256, rotary_percentage=0.5
+        model_name, padded_vocab_size=10000, n_layer=2, n_head=4, n_embd=256, rotary_percentage=0.5
     )
     T = 5
     theirs_config = PhiConfig(