fix pylint

Author: zigzagcai

internlm/checkpoint/components.py

@@ -4,7 +4,6 @@
 from collections import defaultdict
 
 import torch
-from torch.distributed._shard.api import load_with_process_group
 
 from internlm.accelerator import get_accelerator
 from internlm.core.context import ParallelMode
@@ -15,13 +14,10 @@
 from internlm.utils.common import get_current_device
 from internlm.utils.lazy import LazyObject
 from internlm.utils.logger import get_logger
-from internlm.utils.parallel import is_using_hf, is_using_fsdp, is_using_isp
+from internlm.utils.parallel import is_using_fsdp, is_using_hf, is_using_isp
 from internlm.utils.storage_manager import get_fns, llm_load, llm_save
 
-from .utils import (
-    get_model_topology,
-    get_non_moe_state_dict,
-)
+from .utils import get_model_topology, get_non_moe_state_dict
 
 try:
     import torch.distributed.checkpoint as dcp
@@ -194,7 +190,7 @@ def load_model_checkpoint(folder, model):
     else:
         should_load_name = f"model_tp{tp_rank}_pp{pp_rank}.pt"
     fp = os.path.join(folder, should_load_name)
-    
+
     states = llm_load(fp, map_location=get_current_device())
     """
     # need convert the gate parameters to float32 (to fit deepspeed style mechanism), it may cause round-off in
@@ -366,7 +362,7 @@ def load_optimizer_checkpoint(folder, optim):
                     max_pp = max(max_pp, int(pp[2:]))
             else:
                 _, fsdp = os.path.splitext(fn)[0].split("_")
-                max_fsdp = max(max_fsdp, int(fsdp[4:]))  
+                max_fsdp = max(max_fsdp, int(fsdp[4:]))
 
     fsdp_size = gpc.get_world_size(ParallelMode.GLOBAL)
     zero_size = gpc.get_world_size(ParallelMode.ZERO1)
@@ -399,7 +395,7 @@ def load_optimizer_checkpoint(folder, optim):
     tp_rank = gpc.get_local_rank(ParallelMode.TENSOR)
     wp_rank = gpc.get_local_rank(ParallelMode.WEIGHT)
     pp_rank = gpc.get_local_rank(ParallelMode.PIPELINE)
-    
+
     if isinstance(optim, (HybridZeroOptimizer, HybridZeroOptimizer_v2)):
         if is_using_isp():
             fp = f"optimizer_wp{wp_rank}_pp{pp_rank}_zo{zero_rank}.pt"

internlm/checkpoint/utils.py

@@ -2,10 +2,10 @@
 # -*- encoding: utf-8 -*-
 
 import itertools
+
 import numpy as np
 import torch
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-from torch.distributed.fsdp import StateDictType
 
 from internlm.core.context import global_context as gpc
 from internlm.core.parallel.shard import split_data_for_sequence_parallel
@@ -116,4 +116,4 @@ def init_fsdp_v1(model: FSDP, device: torch.device) -> FSDP:
 
         # run a forward pass with dummy_input to initialize FSDP
         _ = model(**dummy_input)
-    return model
+    return model

internlm/core/context/process_group_initializer.py

@@ -1107,7 +1107,6 @@ def init_dist_group(self, use_cpu: bool = False):
         return groups
 
 
-
 class Initializer_Weight(ProcessGroupInitializer):
     """A ProcessGroupInitializer for model weight parallelism.
 

internlm/core/trainer_builder.py

@@ -98,7 +98,7 @@ def __init__(
         self.current_time = self._setup_time_and_logging()
         # load config_lines
         config_lines = self._read_config(kwargs["config"])
-        
+
         # inject model for amp, parallel setting, parameter syncing and others
         model, isp_communicator = inject_model(model)
 

internlm/model/builder.py

@@ -1,19 +1,21 @@
 from typing import List, Union
 
-from torch import nn
 import torch
+from torch import nn
 
 from internlm.core.context import ParallelMode
 from internlm.core.context import global_context as gpc
 from internlm.core.parallel.shard import pipeline_parallel_sharding_wrapper
 from internlm.model.base_model import BaseModel
-from internlm.model.modules.linear import ParallelLinearWithCommExt, ScaleColumnParallelLinear
+from internlm.model.modules.linear import (
+    ParallelLinearWithCommExt,
+    ScaleColumnParallelLinear,
+)
 from internlm.model.registry import model_initializer
-from internlm.utils.parallel import is_using_hf
 from internlm.utils.common import get_current_device
 from internlm.utils.lazy import LazyObject
 from internlm.utils.logger import get_logger
-from internlm.utils.parallel import is_using_fsdp, is_using_isp
+from internlm.utils.parallel import is_using_fsdp, is_using_hf, is_using_isp
 
 logger = get_logger(__file__)
 
@@ -58,6 +60,7 @@ def create_model_builtin(model_type) -> Union[nn.Module, List[nn.Module]]:
 
     return model
 
+
 def create_model_hf(hf: dict) -> nn.Module:
     cfg = LazyObject(hf.cfg, hf.cfg_cls)
     cfg = cfg.build()
@@ -123,4 +126,4 @@ def traverse(module):
         else:
             traverse(model)
 
-    return model
+    return model

internlm/solver/activation_checkpoint.py

@@ -4,11 +4,9 @@
 import weakref
 
 import torch
-
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
     checkpoint_wrapper as ptd_checkpoint_wrapper,
 )
-
 from torch.utils.checkpoint import check_backward_validity, detach_variable
 
 from internlm.accelerator import get_accelerator
@@ -287,4 +285,4 @@ def apply_ac_to_transformer_block(module: torch.nn.Module, checkpoint):
     if ptd_checkpoint_wrapper._count % ac_freq == 0:
         return ptd_checkpoint_wrapper(module, preserve_rng_state=False)
     else:
-        return module
+        return module

internlm/solver/optimizer/fsdp_optimizer.py

@@ -7,7 +7,7 @@
 import torch.distributed as dist
 from torch.optim import Optimizer
 
-from internlm.accelerator import AcceleratorType, get_accelerator
+from internlm.accelerator import get_accelerator
 from internlm.core.context import Config, ParallelMode
 from internlm.core.context import global_context as gpc
 from internlm.solver.optimizer.base_optimizer import BaseOptimizer

internlm/train/pipeline.py

@@ -10,13 +10,13 @@
 
 import torch
 from torch import nn
-from torch.utils.data import DataLoader
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp.fully_sharded_data_parallel import (
     BackwardPrefetch,
     ShardingStrategy,
 )
 from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
+from torch.utils.data import DataLoader
 
 from internlm.accelerator import AcceleratorType, get_accelerator
 from internlm.checkpoint.utils import init_fsdp_v1
@@ -96,8 +96,8 @@
     is_replica_zero_parallel_parameter,
     is_tensor_expert_data_parallel_parameter,
     is_tensor_zero_parallel_parameter,
-    is_using_hf,
     is_using_fsdp,
+    is_using_hf,
     is_using_isp,
     is_weight_expert_data_parallel_parameter,
     is_weight_zero_parallel_parameter,
@@ -256,7 +256,8 @@ def _check_module(name, module):
             # special case for pure dp mode
             if (
                 isinstance(gpc.config.parallel["tensor"], dict)
-                and gpc.config.parallel["tensor"].get("mode", TensorParallelMode.mtp.name) == TensorParallelMode.mtp.name
+                and gpc.config.parallel["tensor"].get("mode", TensorParallelMode.mtp.name)
+                == TensorParallelMode.mtp.name
                 and gpc.get_world_size(ParallelMode.DATA) == gpc.get_world_size(ParallelMode.GLOBAL)
             ):
                 _check_module_func = _check_module_pure_dp
@@ -278,7 +279,9 @@ def _check_module(name, module):
 
 
 @llm_timeout(func_name="initialize_model_and_parallel_communicator")
-def initialize_model_and_parallel_communicator(pre_process_func: Optional[Callable] = None, post_process_func: Optional[Callable] = None):
+def initialize_model_and_parallel_communicator(
+    pre_process_func: Optional[Callable] = None, post_process_func: Optional[Callable] = None
+):
     """
     Initialize model with Automatic Mixed Precision.
     Returns:
@@ -362,10 +365,10 @@ def inject_model(model):
     # state in the same dp group are all the same.
     random_mode = ParallelMode.WEIGHT_DATA if is_using_isp() else ParallelMode.DATA
     set_mode(random_mode)
-    
+
     # initialize isp communicator
     isp_communicator = initialize_parallel_communicator(model)
-    
+
     model = wrap_FSDP_model(model)
 
     # set is_injected flag

internlm/utils/parallel.py

@@ -27,6 +27,7 @@ def is_using_fsdp():
         and gpc.config.parallel["fsdp"].get("enable", False)
     )
 
+
 def is_using_sequence_parallel():
     return (
         isinstance(gpc.config.parallel["tensor"], dict)