Apply isort and black reformatting

Signed-off-by: kevalmorabia97 <kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

nemo/collections/llm/modelopt/distill/utils.py

@@ -39,7 +39,7 @@
     from nemo.collections.common.tokenizers.tokenizer_spec import TokenizerSpec
 
 import modelopt.torch.opt as mto
-from modelopt.torch.distill import DistillationModel, DistillationLossBalancer
+from modelopt.torch.distill import DistillationLossBalancer, DistillationModel
 
 
 @dataclass

nemo/collections/llm/modelopt/model_utils.py

@@ -18,6 +18,7 @@
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import lightning.pytorch as L
+import modelopt.torch.opt as mto
 import torch
 import torch.nn as nn
 from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
@@ -32,8 +33,6 @@
 from nemo.utils.import_utils import safe_import
 from nemo.utils.model_utils import unwrap_model
 
-import modelopt.torch.opt as mto
-
 _, HAVE_TE = safe_import("transformer_engine")
 if HAVE_TE:
     # These custom modelopt specs are a mix of local MCORE and TE specs.

nemo/collections/llm/modelopt/prune/pruner.py

@@ -15,10 +15,10 @@
 from dataclasses import dataclass
 from functools import partial
 
+import modelopt.torch.prune as mtp
 import pytorch_lightning as pl
 from megatron.core import dist_checkpointing
 
-import modelopt.torch.prune as mtp
 from nemo import lightning as nl
 from nemo.collections import llm
 from nemo.lightning.ckpt_utils import ckpt_to_context_subdir

nemo/collections/llm/modelopt/quantization/quantizer.py

@@ -21,6 +21,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional, Union
 
+import modelopt.torch.export as mte
+import modelopt.torch.opt as mto
+import modelopt.torch.quantization as mtq
 import torch
 from datasets import load_dataset
 from megatron.core.inference.common_inference_params import CommonInferenceParams
@@ -39,18 +42,13 @@
 from nemo.utils.get_rank import is_global_rank_zero
 from nemo.utils.model_utils import unwrap_model
 
-import modelopt.torch.export as mte
-import modelopt.torch.quantization as mtq
-import modelopt.torch.opt as mto
-
 if TYPE_CHECKING:
     import lightning.pytorch as pl
 
     from nemo.lightning import Trainer
     from nemo.lightning.megatron_parallel import MegatronParallel
 
 
-
 QUANT_CFG_CHOICES = get_quant_cfg_choices()
 SUPPORTED_DTYPE = [16, "16", "bf16"]  # Default precision for non-quantized layers
 SUPPORTED_EXPORT_FMT = ["trtllm", "nemo", "hf"]

nemo/collections/llm/modelopt/speculative/model_transform.py

@@ -12,10 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch.nn as nn
-
 import modelopt.torch.opt as mto
 import modelopt.torch.speculative as mtsp
+import torch.nn as nn
+
 from nemo.collections.llm import GPTModel
 from nemo.utils import logging
 from nemo.utils.model_utils import unwrap_model

nemo/export/quantize/quantizer.py

@@ -17,10 +17,12 @@
 from contextlib import nullcontext
 from typing import Callable, Optional
 
+import modelopt.torch.quantization as mtq
 import torch
 import torch.distributed as dist
 from megatron.core import parallel_state
 from megatron.core.transformer.module import Float16Module
+from modelopt.torch.export import export_tensorrt_llm_checkpoint
 from omegaconf.omegaconf import DictConfig, open_dict
 
 from nemo.collections.nlp.models.language_modeling.megatron_gpt_model import MegatronGPTModel
@@ -29,9 +31,6 @@
 from nemo.utils.distributed import temporary_directory
 from nemo.utils.model_utils import save_artifacts, unwrap_model
 
-import modelopt.torch.quantization as mtq
-from modelopt.torch.export import export_tensorrt_llm_checkpoint
-
 QUANT_CFG_CHOICES = {
     "int8": mtq.INT8_DEFAULT_CFG,
     "int8_sq": mtq.INT8_SMOOTHQUANT_CFG,

nemo/lightning/fabric/strategies.py

@@ -28,6 +28,7 @@
     Union,
 )
 
+import modelopt.torch.opt as mto
 import torch
 from lightning.fabric.accelerators import CPUAccelerator
 from lightning.fabric.accelerators.accelerator import Accelerator
@@ -60,8 +61,6 @@
 from nemo.utils.import_utils import safe_import
 from nemo.utils.model_utils import unwrap_model
 
-import modelopt.torch.opt as mto
-
 if TYPE_CHECKING:
     from nemo.lightning.pytorch.plugins.data_sampler import DataSampler
 

scripts/llm/gpt_prune.py

@@ -66,6 +66,7 @@
 # isort: off
 # Import modelopt first to avoid circular import in 0.35.0
 import modelopt.torch.prune  # noqa: F401
+
 # isort: on
 
 from nemo.collections import llm