[Core] Fix import error (#386)

To accommodate the scenario where the before_start field is used to
switch to the actual environment during program execution, we have
placed the import statements inside the function body rather than at the
beginning of the file.

Co-authored-by: caozhou <[email protected]>

Author: Caozhou1995

flagscale/runner/auto_tuner/memory_model.py

@@ -1,9 +1,10 @@
 from flagscale.runner.auto_tuner.utils import convert_config_to_megatron_args
-from flagscale.train.theoretical_memory_usage import report_theoretical_memory
 
 
 def default_model(strategy, config):
     """Use megatron built in memory model."""
+    from flagscale.train.theoretical_memory_usage import report_theoretical_memory
+
     args = convert_config_to_megatron_args(config, strategy)
     num_microbatches = (
         config.train.model.global_batch_size

run.py

@@ -8,6 +8,9 @@
 from flagscale.runner.runner_train import CloudTrainRunner, SSHTrainRunner
 from flagscale.runner.utils import is_master
 
+# To accommodate the scenario where the before_start field is used to switch to the actual environment during program execution,
+# we have placed the import statements inside the function body rather than at the beginning of the file.
+
 
 @hydra.main(version_base=None, config_name="config")
 def main(config: DictConfig) -> None: