Skip to content

GitLab

Commit d58fa705 authored by Yuanchen's avatar Yuanchen Committed by binmakeswell
Browse files

[NFC] polish code style (#3268) 


Co-authored-by: default avatarYuanchen Xu <yuanchen.xu00@gmail.com>
parent c4a226b7
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 6 deletions
colossalai/nn/_ops/_utils.py
View file @ d58fa705
import torch from typing import List, Optional, Union
from typing import Union, Optional, List
from colossalai.tensor import ColoTensor
import torch import torch
import torch.distributed as dist import torch.distributed as dist
from colossalai.global_variables import tensor_parallel_env as env
from colossalai.global_variables import tensor_parallel_env as env
from colossalai.nn.layer.utils import divide from colossalai.nn.layer.utils import divide
from colossalai.tensor import ProcessGroup, ColoTensorSpec from colossalai.tensor import ColoTensor, ColoTensorSpec, ProcessGroup
GeneralTensor = Union[ColoTensor, torch.Tensor] GeneralTensor = Union[ColoTensor, torch.Tensor]
Number = Union[int, float] Number = Union[int, float]
...@@ -135,7 +134,7 @@ class _ReduceInput(torch.autograd.Function): ...@@ -135,7 +134,7 @@ class _ReduceInput(torch.autograd.Function):
class _SplitForwardGatherBackward(torch.autograd.Function): class _SplitForwardGatherBackward(torch.autograd.Function):
""" """
Split the input and keep only the corresponding chuck to the rank. Split the input and keep only the corresponding chuck to the rank.
Args: Args:
input_: input matrix. input_: input matrix.
process_group: parallel mode. process_group: parallel mode.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment