github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-17 07:26:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish colossalai/cli/benchmark/benchmark.py code style (#2287)

Browse Source
This commit is contained in:
アマデウス 2023-01-03 18:23:11 +08:00 committed by Frank Lee
parent 1c29b173c9
commit 49715a78f0
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
colossalai/cli/benchmark/benchmark.py
View File

@ -1,16 +1,17 @@
import colossalai from functools import partial
from typing import Dict, List
import click import click
import torch.multiprocessing as mp import torch.multiprocessing as mp
from functools import partial import colossalai
from typing import List, Dict from colossalai.cli.benchmark.utils import find_all_configs, get_batch_data, profile_model
from colossalai.context import Config from colossalai.context import Config
from colossalai.context.random import reset_seeds from colossalai.context.random import reset_seeds
from colossalai.core import global_context as gpc from colossalai.core import global_context as gpc
from colossalai.logging import disable_existing_loggers, get_dist_logger from colossalai.logging import disable_existing_loggers, get_dist_logger
from colossalai.utils import free_port, MultiTimer from colossalai.utils import MultiTimer, free_port
from colossalai.cli.benchmark.utils import find_all_configs, profile_model, get_batch_data
from .models import MLP from .models import MLP
@ -53,7 +54,7 @@ def run_dist_profiling(rank: int, world_size: int, port_list: List[int], config_
port_list (List[int]): a list of free ports for initializing distributed networks port_list (List[int]): a list of free ports for initializing distributed networks
config_list (List[Dict]): a list of configuration config_list (List[Dict]): a list of configuration
hyperparams (Config): the hyperparameters given by the user hyperparams (Config): the hyperparameters given by the user
""" """
# disable logging for clean output # disable logging for clean output