flux bug

dcu_megatron/adaptor/megatron_adaptor.py

@@ -188,6 +188,13 @@ class CoreAdaptation(MegatronAdaptationABC):
                                     apply_wrapper=True)
 
         # flux
+        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)

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