Repository: pytorch/torchtune

Issue:
LORAPhi3 bug when merging lora weights 
I noticed a strange behavior that only happens to LoraPhi3 model (everything works fine for llama or gemma): after merging weights back to non Lora model the non Lora model output results completely different than the Lora version.

I have not been able to spot where the problem is coming from exactly but I can share a minimal reproducible code to illustrate what I consider a bug:

```
from torchtune.models.phi3 import lora_phi3, phi3
from torchtune.modules.peft.peft_utils import get_merged_lora_ckpt
import torch

# let's create a small 4 layers lora_phi3
lora_backbone = lora_phi3(
                                            lora_attn_modules=['q_proj', 'v_proj', 'k_proj', 'output_proj'],
                                            apply_lora_to_mlp=False,
                                            apply_lora_to_output=False,
                                            vocab_size=32_064,
                                            num_layers=4,
                                            num_heads=32,
                                            num_kv_heads=32,
                                            embed_dim=3072,
                                            intermediate_dim=8192,
                                            max_seq_len=4096,
                                            attn_dropout=0.,
                                            norm_eps=1e-5,
                                            lora_rank=8,
                                            lora_alpha=16,
                                            lora_dropout=0.,
                                            quantize_base=False,
                                        )
# now let's simply merge the lora weights back to the original phi3 model
merged_state_dict = get_merged_lora_ckpt(
                lora_backbone.state_dict(),
                rank=8,
                alpha=16,
            )

# create the corresponding original phi3 model and load the merged weights
backbone = phi3(
                                            
                                            vocab_size=32_064,
                                            num_layers=4,
                                            num_heads=32,
                                            num_kv_heads=32,
                                            embed_dim=3072,
                                            intermediate_dim=8192,
                                            max_seq_len=4096,
                                            attn_dropout=0.,
                                            norm_eps=1e-5,
                                        )
backbone.load_state_dict(merged_state_dict)

# now let's compare the different in outputs on a random set of tokens
a = torch.randint(0, 32_000, (1, 128)).long()
with torch.no_grad():
    lora_output = lora_backbone(a)
    output = backbone(a)

# This returns False!
torch.allclose(lora_output, output, atol=1e-01)

```

This is obviously just a toy example but in real applications the predictions are almost completely random so there is a real problem here.
I'll try to open a PR if I do spot what is going wrong here!


Discussion:
Looking into this now, thanks for the pointer!

cc @ebsmothers 
@joecummings found it! If you add the same rope_base when defining both models then everythings works as expected, the problem comes from the fact that the default `robe_base`  values are different in lora and non lora version.

I have proposed a simple fix here : #997
I think 10000.0 matches hugging face values : https://huggingface.co/microsoft/Phi-3-mini-128k-instruct/blob/main/config.json
Assuming that `rope_theta` from hugging face corresponds to `rope_base` from torchtune.



Patch:
diff --git a/torchtune/models/phi3/_component_builders.py b/torchtune/models/phi3/_component_builders.py
index c015380a60..da58dcf39d 100644
--- a/torchtune/models/phi3/_component_builders.py
+++ b/torchtune/models/phi3/_component_builders.py
@@ -133,7 +133,7 @@ def lora_phi3(
     max_seq_len: int,
     attn_dropout: float = 0.0,
     norm_eps: float = 1e-5,
-    rope_base: float = 500000.0,
+    rope_base: int = 10_000,
     # LoRA args
     lora_rank: int,
     lora_alpha: float,
@@ -167,6 +167,8 @@ def lora_phi3(
         attn_dropout (float): dropout value passed onto scaled_dot_product_attention.
             Default: 0.0
         norm_eps (float): epsilon in RMS norms.
+        rope_base (int): base value for Rotary Position Embeddings.
+            Default: 10000
         lora_rank (int): rank of each low-rank approximation
         lora_alpha (float): scaling factor for the low-rank approximation
         lora_dropout (float): LoRA dropout probability. Default: 0.0
@@ -252,7 +254,7 @@ def lora_phi3_self_attention(
     num_kv_heads: int,
     max_seq_len: int,
     attn_dropout: float = 0.0,
-    rope_base: float = 500000.0,
+    rope_base: int = 10_000,
     # LoRA args
     lora_rank: int,
     lora_alpha: float,
@@ -277,6 +279,8 @@ def lora_phi3_self_attention(
             by :func:`~torchtune.modules.KVCache`
         attn_dropout (float): dropout value passed onto scaled_dot_product_attention.
             Default: 0.0
+        rope_base (int): base value for Rotary Position Embeddings.
+            Default: 10000
         lora_rank (int): rank of each low-rank approximation
         lora_alpha (float): scaling factor for the low-rank approximation
         lora_dropout (float): LoRA dropout probability. Default: 0.0


Test patch:
diff --git a/tests/torchtune/models/phi3/test_lora_phi3.py b/tests/torchtune/models/phi3/test_lora_phi3.py
index cef4dbcd2c..148e9c6456 100644
--- a/tests/torchtune/models/phi3/test_lora_phi3.py
+++ b/tests/torchtune/models/phi3/test_lora_phi3.py
@@ -61,9 +61,9 @@ def test_empty_lora_modules(self):
     @pytest.mark.parametrize(
         "lora_modules, expected",
         [
-            (["q_proj", "v_proj"], torch.tensor(50.64567)),
-            (["q_proj", "k_proj", "v_proj", "output_proj"], torch.tensor(77.17097)),
-            (["k_proj"], torch.tensor(44.90090)),
+            (["q_proj", "v_proj"], torch.tensor(51.88655)),
+            (["q_proj", "k_proj", "v_proj", "output_proj"], torch.tensor(75.80934)),
+            (["k_proj"], torch.tensor(44.00425)),
         ],
     )
     def test_forward(self, inputs, lora_modules, expected):
@@ -137,14 +137,14 @@ def get_ref_phi3(self, vocab_size, embed_dim=EMBED_DIM):
     @pytest.mark.parametrize(
         "lora_modules, apply_lora_to_mlp, apply_lora_to_output, expected",
         [
-            (["q_proj", "v_proj"], False, False, torch.tensor(2869701.75)),
+            (["q_proj", "v_proj"], False, False, torch.tensor(2869687.75)),
             (
                 ["q_proj", "k_proj", "v_proj", "output_proj"],
                 True,
                 False,
-                torch.tensor(10674807.0),
+                torch.tensor(10674772.0),
             ),
-            (["k_proj"], True, True, torch.tensor(16285855.0)),
+            (["k_proj"], True, True, torch.tensor(16285834.0)),
         ],
     )
     def test_forward(