fix

Author: tastelikefeet

src/twinkle/model/megatron/megatron.py

@@ -34,8 +34,6 @@
 from twinkle.template import Template
 from .strategy import MegatronStrategy
 from twinkle.utils import construct_class, exists
-from .args import get_args, set_args, TwinkleMegatronArgs
-from .model import get_megatron_model_meta, GPTBridge
 from twinkle.patch import Patch, apply_patch
 
 
@@ -173,6 +171,7 @@ def __init__(
         **kwargs,
     ):
         requires('megatron_core')
+        from .args import get_args, set_args, TwinkleMegatronArgs
         os.environ['TOKENIZERS_PARALLELISM'] = 'true'
         nn.Module.__init__(self)
         from twinkle.patch.megatron_peft import MegatronPeft
@@ -240,6 +239,7 @@ def _create_megatron_model(
         load_weights: bool = True,
         **kwargs,
     ) -> List[nn.Module]:
+        from .args import get_args
         args = get_args()
         self.initialize(**kwargs)
 
@@ -1002,6 +1002,7 @@ def _patch_adapter(self, adapter_name: str, config_or_dir: Union[PeftConfig, str
                     if isinstance(m, LoraLinear):
                         # just check
                         # TODO untested code
+                        from .args import get_args
                         args = get_args()
                         from .tuners import LoraParallelLinear
                         assert args.is_multimodal and not isinstance(m, LoraParallelLinear)
@@ -1114,6 +1115,7 @@ def initialize(self, **kwargs) -> None:
 
         from megatron.core import parallel_state
         from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
+        from .args import get_args
         self._try_init_process_group()
         args = get_args()
         init_kwargs = {
@@ -1142,8 +1144,10 @@ def initialize(self, **kwargs) -> None:
         self._initialized = True
 
     @property
-    def _bridge(self) -> GPTBridge:
+    def _bridge(self) -> 'GPTBridge':
         if not hasattr(self, '_bridge_instance'):
+            from .args import get_args
+            from .model import get_megatron_model_meta
             args = get_args()
             megatron_model_meta = get_megatron_model_meta(args.hf_model_type)
             assert megatron_model_meta is not None, f'Model: {args.hf_model_type} is not supported.'
@@ -1181,6 +1185,7 @@ def send_weights(
         # Trim any tensor whose dim-0 equals padded_vocab_size back to
         # org_vocab_size — this is shape-based, not name-based, so it works
         # regardless of the model architecture's naming convention.
+        from .args import get_args
         args = get_args()
         org_vocab_size = getattr(self.hf_config, 'vocab_size', args.padded_vocab_size)
         _padded_vocab_size = args.padded_vocab_size

src/twinkle/model/megatron/multi_lora_megatron.py

@@ -17,7 +17,6 @@
 from twinkle.loss import Loss
 from twinkle.metric import Metric
 from twinkle.processor import InputProcessor
-from .args import TwinkleMegatronArgs, set_args
 from .megatron import MegatronModel
 from .strategy import MegatronStrategy
 from ..multi_lora import MultiLora
@@ -42,6 +41,7 @@ def __init__(self,
         requires('megatron_core')
         os.environ['TOKENIZERS_PARALLELISM'] = 'true'
         os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1"
+        from .args import TwinkleMegatronArgs, set_args
         nn.Module.__init__(self)
         from twinkle.patch.megatron_peft import MegatronPeft