Manual pre-commit lint.

Signed-off-by: Cory Ye <cye@nvidia.com>

Author: cspades

bionemo-recipes/recipes/esm2_native_te/train_mfsdp.py

@@ -20,6 +20,7 @@
 import torch
 import transformer_engine.pytorch
 import transformers
+from megatron_fsdp.fully_shard import fully_shard
 from omegaconf import DictConfig, OmegaConf
 from torch.distributed.device_mesh import init_device_mesh
 from torch.optim import AdamW
@@ -30,7 +31,6 @@
 from checkpoint import load_checkpoint_mfsdp, save_checkpoint_mfsdp, save_final_model_mfsdp, should_save_checkpoint
 from dataset import create_bshd_dataloader, create_thd_dataloader
 from distributed_config import DistributedConfig
-from megatron_fsdp.fully_shard import fully_shard
 from modeling_esm_te import NVEsmConfig, NVEsmForMaskedLM
 from perf_logger import PerfLogger
 from quantization import resolve_layer_precision

bionemo-recipes/recipes/vit/train.py

@@ -22,6 +22,7 @@
 import torch
 import wandb
 from hydra.core.hydra_config import HydraConfig
+from megatron_fsdp.fully_shard import fully_shard
 from omegaconf import OmegaConf
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import CosineAnnealingLR
@@ -31,7 +32,6 @@
 from beans import BeansDataset, infinite_dataloader
 from checkpoint import load_auto_resume_checkpoint, save_auto_resumable_checkpoint
 from distributed import initialize_distributed
-from megatron_fsdp import fully_shard
 from vit import build_vit_model