fix flux import error

dcu_megatron/adaptor/megatron_adaptor.py

@@ -165,7 +165,7 @@ class CoreAdaptation(MegatronAdaptationABC):
 
         # flux
         if int(os.getenv("USE_FLUX_OVERLAP", "0")):
-            from ..core.tensor_parallel import (
+            from ..core.tensor_parallel.layers import (
                 FluxColumnParallelLinear,
                 FluxRowParallelLinear
             )

dcu_megatron/core/models/gpt/gpt_model.py

@@ -12,8 +12,6 @@ from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.utils import WrappedTensor, deprecate_inference_params
 
-from dcu_megatron.core.tensor_parallel import FluxColumnParallelLinear
-
 
 def gpt_model_init_wrapper(fn):
     @wraps(fn)
@@ -25,6 +23,8 @@ def gpt_model_init_wrapper(fn):
             (self.post_process or self.mtp_process)
             and int(os.getenv("USE_FLUX_OVERLAP", "0"))
         ):
+            from dcu_megatron.core.tensor_parallel.layers import FluxColumnParallelLinear
+
             self.output_layer = FluxColumnParallelLinear(
                 self.config.hidden_size,
                 self.vocab_size,

dcu_megatron/core/tensor_parallel/__init__.py

-from .layers import (
-    FluxColumnParallelLinear,
-    FluxRowParallelLinear,
-)
\ No newline at end of file