fix(cli): remove import to avoid importing neuronx_distributed

Author: tengomucho

optimum/commands/neuron/subcommands.py

@@ -16,7 +16,6 @@
 
 from typing import TYPE_CHECKING
 
-from ...neuron.models.training import consolidate_model_parallel_checkpoints_to_unified_checkpoint
 from ...utils import logging
 from ..base import BaseOptimumCLICommand
 
@@ -53,6 +52,9 @@ def parse_args(parser: "ArgumentParser"):
         )
 
     def run(self):
+        # This is not on top otherwise it will make the CLI require neuronx_distributed
+        from ...neuron.models.training import consolidate_model_parallel_checkpoints_to_unified_checkpoint
+
         checkpoint_format = "safetensors" if self.args.format == "safetensors" else "pytorch"
         logger.info(f"Consolidating checkpoints from {self.args.checkpoint_dir} to the {checkpoint_format} format...")
         output_dir = self.args.output_dir