fix

Author: tastelikefeet

src/twinkle/model/megatron/utils/utils.py

@@ -3,15 +3,12 @@
 """
 from typing import Optional
 import torch
-from twinkle import exists, requires
+from twinkle import requires
 
-mpu = None
-if exists('megatron_core'):
-    from megatron.core import mpu  # type: ignore
 
 def split_cp_inputs(inputs: torch.Tensor, cu_seqlens: Optional[torch.Tensor], dim: int):
-    if mpu is None:
-        requires('megatron_core')
+    requires('megatron_core')
+    from megatron.core import mpu
     if dim < 0:
         dim = (dim + inputs.ndim) % inputs.ndim
     new_inputs = []