test: fix neuronx_distributed import

huggingface · Jan 2, 2025 · a2039d4 · a2039d4
1 parent 4cd2843
commit a2039d4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/distributed/test_model_parallelization.py b/tests/distributed/test_model_parallelization.py
@@ -64,8 +64,8 @@
     import torch_xla.core.xla_model as xm
 
 if is_neuronx_distributed_available():
-    from neuronx_distributed.modules.qkv_linear import get_kv_shared_group
     from neuronx_distributed.parallel_layers.parallel_state import (
+        get_kv_shared_group,
         get_pipeline_model_parallel_rank,
         get_tensor_model_parallel_group,
         get_tensor_model_parallel_size,