Fixed test

optimum/neuron/distributed/utils.py

@@ -1020,9 +1020,8 @@ def _parallelize_active_adapters(
     skip_weight_load: bool = False,
     device: Optional["torch.device"] = None,
 ):
-    from peft.tuners.lora import Linear as LoraLinear
     from neuronx_distributed.parallel_layers.parallel_state import get_tensor_model_parallel_size
-
+    from peft.tuners.lora import Linear as LoraLinear
 
     try:
         peft_config = tuner_layer._peft_config
@@ -1051,9 +1050,11 @@ def _parallelize_active_adapters(
                 dim_to_partition = layer_to_parallelize.weight.size(1)
 
             tp_size = get_tensor_model_parallel_size()
-            
+
             if dim_to_partition % tp_size != 0:
-                raise RuntimeError(f"The LoRA adapter dimension to parallelize ({dim_to_partition}) is not divisible by the TP size ({tp_size}).")
+                raise RuntimeError(
+                    f"The LoRA adapter dimension to parallelize ({dim_to_partition}) is not divisible by the TP size ({tp_size})."
+                )
 
             # TODO: handle the case were weights already exist for this adapter.
             parallel_layer = linear_to_parallel_linear(

tests/distributed/test_model_parallelization.py

@@ -23,7 +23,7 @@
 import torch.utils._pytree as pytree
 from peft import LoraConfig
 from peft import get_peft_model as orig_get_peft_model
-from transformers import AutoTokenizer, LlamaForCausalLM, AutoModelForCausalLM, AutoConfig
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, LlamaForCausalLM
 from transformers.models.auto.configuration_auto import CONFIG_MAPPING
 from transformers.models.auto.modeling_auto import (
     MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING,
@@ -429,17 +429,23 @@ def test_parallel_model_matches_original_model_from_config(
         )
 
     @pytest.mark.parallel_sizes((8, 8, 1))
-    def test_llama_v2_gqa_and_lora(self, lazy_load):
+    def test_llama_v2_gqa_and_lora(self, lazy_load, tmpdir):
         tp_size = get_tensor_model_parallel_size()
         tp_group = get_tensor_model_parallel_group()
 
         static_seed_patcher = StaticSeedPatcher(42)
 
+        # First we create and save a modified version of the model that is suited for our test case.
         config = AutoConfig.from_pretrained(LLAMA_V2_MODEL_NAME)
         config.num_hidden_layers = 2
         config.num_key_value_heads = 2
         config.hidden_size = 128
-        config.tie_word_embeddings = True
+        modified_model = AutoModelForCausalLM.from_pretrained(
+            LLAMA_V2_MODEL_NAME, config=config, ignore_mismatched_sizes=True
+        )
+        if xm.get_ordinal() == 0:
+            modified_model.save_pretrained(tmpdir)
+        xm.rendezvous("Saved modified model")
 
         lora_config = LoraConfig(
             r=64,
@@ -449,9 +455,8 @@ def test_llama_v2_gqa_and_lora(self, lazy_load):
             task_type="CAUSAL_LM",
         )
 
-        with static_seed_patcher:
-            orig_model = AutoModelForCausalLM.from_pretrained(LLAMA_V2_MODEL_NAME, config=config, ignore_mismatched_sizes=True)
-            orig_model.eval()
+        orig_model = AutoModelForCausalLM.from_pretrained(tmpdir)
+        orig_model.eval()
 
         with static_seed_patcher:
             orig_model = orig_get_peft_model(orig_model, lora_config)
@@ -460,10 +465,9 @@ def test_llama_v2_gqa_and_lora(self, lazy_load):
 
         ctx = lazy_load_for_parallelism(tensor_parallel_size=tp_size) if lazy_load else nullcontext()
         with ctx:
-            with static_seed_patcher:
-                model = AutoModelForCausalLM.from_pretrained(LLAMA_V2_MODEL_NAME, config=config, ignore_mismatched_sizes=True)
-                model.eval()
-
+            model = AutoModelForCausalLM.from_pretrained(tmpdir)
+            model.eval()
+
         with static_seed_patcher:
             model = get_peft_model(model, lora_config)
 
@@ -480,7 +484,6 @@ def test_llama_v2_gqa_and_lora(self, lazy_load):
         tok.pad_token = tok.eos_token
 
         inputs = tok("This is a curious test.", padding="max_length", max_length=24, return_tensors="pt")
-        print(inputs["input_ids"])
         inputs = {k: v.to("xla") for k, v in inputs.items()}
 
         orig_logits = orig_model(**inputs).logits
@@ -493,7 +496,6 @@ def test_llama_v2_gqa_and_lora(self, lazy_load):
         torch.distributed.all_gather(gathered, logits, group=tp_group)
         gathered_logits = torch.cat(gathered, dim=2)
         xm.mark_step()
-        xm.master_print(torch.nonzero(orig_logits - gathered_logits))
         torch.testing.assert_close(orig_logits, gathered_logits)
 
     @pytest.mark.skipif(