Skip to content

Commit

Permalink
Apply isort and black reformatting
Browse files Browse the repository at this point in the history
Signed-off-by: JimmyZhang12 <JimmyZhang12@users.noreply.github.com>
  • Loading branch information
JimmyZhang12 committed Sep 21, 2024
1 parent 5da9ee0 commit e941b25
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions nemo/lightning/pytorch/strategies/megatron_strategy.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import shutil
from collections import OrderedDict
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass, asdict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import (
TYPE_CHECKING,
Expand All @@ -40,10 +40,9 @@
import torch.distributed
from lightning_fabric.plugins import CheckpointIO, ClusterEnvironment
from lightning_fabric.utilities.optimizer import _optimizer_to_device, _optimizers_to_device
from megatron.core.dist_checkpointing.core import maybe_load_config
from megatron.core.distributed import DistributedDataParallelConfig
from megatron.core.optimizer import OptimizerConfig
from megatron.core.dist_checkpointing.core import maybe_load_config

from pytorch_lightning.accelerators import CPUAccelerator
from pytorch_lightning.loops import _AutomaticOptimization, evaluation_loop, fit_loop, prediction_loop
from pytorch_lightning.loops.fetchers import _DataLoaderIterDataFetcher
Expand Down Expand Up @@ -674,7 +673,9 @@ def load_checkpoint(self, checkpoint_path: Union[str, Path], selective_restore:
and self.trainer.state.fn == TrainerFn.FITTING
):
if self.lightning_module.optimizers(use_pl_optimizer=False):
sharded_state_dict["optimizer"] = [self.optimizer_sharded_state_dict(is_loading=True, metadata=metadata)]
sharded_state_dict["optimizer"] = [
self.optimizer_sharded_state_dict(is_loading=True, metadata=metadata)
]

checkpoint = self.checkpoint_io.load_checkpoint(checkpoint_path, sharded_state_dict=sharded_state_dict)

Expand Down

0 comments on commit e941b25

Please sign in to comment.