flux bug

dcu_megatron/adaptor/megatron_adaptor.py

@@ -188,16 +188,23 @@ class CoreAdaptation(MegatronAdaptationABC):
                                     apply_wrapper=True)
 
         # flux
-        MegatronAdaptation.register("megatron.core.tensor_parallel.layers.ColumnParallelLinear.__init__",
-                                    parallel_linear_init_wrapper,
-                                    apply_wrapper=True)
-        MegatronAdaptation.register("megatron.core.tensor_parallel.layers.ColumnParallelLinear.forward",
-                                    ColumnParallelLinearPatch.forward)
-        MegatronAdaptation.register("megatron.core.tensor_parallel.layers.RowParallelLinear.__init__",
-                                    parallel_linear_init_wrapper,
-                                    apply_wrapper=True)
-        MegatronAdaptation.register("megatron.core.tensor_parallel.layers.RowParallelLinear.forward",
-                                    RowParallelLinearPatch.forward)
+        try:
+            import flux
+            HAS_FLUX = True
+        except ImportError:
+            HAS_FLUX = False
+
+        if HAS_FLUX:
+            MegatronAdaptation.register("megatron.core.tensor_parallel.layers.ColumnParallelLinear.__init__",
+                                        parallel_linear_init_wrapper,
+                                        apply_wrapper=True)
+            MegatronAdaptation.register("megatron.core.tensor_parallel.layers.ColumnParallelLinear.forward",
+                                        ColumnParallelLinearPatch.forward)
+            MegatronAdaptation.register("megatron.core.tensor_parallel.layers.RowParallelLinear.__init__",
+                                        parallel_linear_init_wrapper,
+                                        apply_wrapper=True)
+            MegatronAdaptation.register("megatron.core.tensor_parallel.layers.RowParallelLinear.forward",
+                                        RowParallelLinearPatch.forward)
 
 
     def patch_training(self):

dcu_megatron/core/tensor_parallel/layers.py

@@ -3,7 +3,12 @@ import warnings
 from functools import wraps
 from typing import Callable, List, Optional
 
-import flux
+try:
+    import flux
+except ImportError:
+    from megatron.training import print_rank_0
+    print_rank_0(f"flux is NOT installed")
+
 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter