github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-25 06:52:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

html refactor (#555)

Browse Source
This commit is contained in:
Liang Bowen 2022-03-31 11:36:56 +08:00 committed by GitHub
parent d1211148a7
commit 2c45efc398
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
133 changed files with 274 additions and 672 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/amp/__init__.py
View File

@ -19,7 +19,7 @@ def convert_to_amp(model: nn.Module, optimizer: Optimizer, criterion: _Loss, mod
optimizer (:class:`torch.optim.Optimizer`): your optimizer object. optimizer (:class:`torch.optim.Optimizer`): your optimizer object.
criterion (:class:`torch.nn.modules.loss._Loss`): your loss function object. criterion (:class:`torch.nn.modules.loss._Loss`): your loss function object.
mode (:class:`colossalai.amp.AMP_TYPE`): amp mode. mode (:class:`colossalai.amp.AMP_TYPE`): amp mode.
amp_config (:class:`colossalai.context.Config` or dict): configuration for different amp modes amp_config (Union[:class:`colossalai.context.Config`, dict]): configuration for different amp modes.
Returns: Returns:
A tuple (model, optimizer, criterion). A tuple (model, optimizer, criterion).

2
colossalai/amp/apex_amp/__init__.py
View File

@ -9,7 +9,7 @@ def convert_to_apex_amp(model: nn.Module, optimizer: Optimizer, amp_config):
Args: Args:
model (:class:`torch.nn.Module`): your model object. model (:class:`torch.nn.Module`): your model object.
optimizer (:class:`torch.optim.Optimizer`): your optimizer object. optimizer (:class:`torch.optim.Optimizer`): your optimizer object.
amp_config (:class: colossalai.context.Config or dict): configuration for initializing apex_amp. amp_config (Union[:class:`colossalai.context.Config`, dict]): configuration for initializing apex_amp.
The ``amp_config`` should include parameters below: The ``amp_config`` should include parameters below:
:: ::

15
colossalai/builder/builder.py
View File

@ -29,21 +29,22 @@ def build_from_registry(config, registry: Registry):
is specified by `registry`. is specified by `registry`.
Note: Note:
the `config` is used to construct the return object such as `LAYERS`, the `config` is used to construct the return object such as `LAYERS`, `OPTIMIZERS`
`OPTIMIZERS` and other support types in `registry`. The `config` should contain and other support types in `registry`. The `config` should contain
all required parameters of corresponding object. The details of support all required parameters of corresponding object. The details of support
types in `registry` and the `mod_type` in `config` could be found in types in `registry` and the `mod_type` in `config` could be found in
`registry <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/registry/__init__.py>`_. `registry <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/registry/__init__.py>`_.
Args: Args:
config (dict or :class:`colossalai.context.colossalai.context.Config`): information config (dict or :class:`colossalai.context.colossalai.context.Config`): information
used in the construction of the return object. used in the construction of the return object.
registry (:class:`Registry`): A registry specifying the type of the return object registry (:class:`Registry`): A registry specifying the type of the return object
Returns: A Python object specified by `registry` Returns:
A Python object specified by `registry`.
Raises: Raises:
Exception: Raises an Exception if an error occurred when building from registry Exception: Raises an Exception if an error occurred when building from registry.
""" """
config_ = config.copy() # keep the original config untouched config_ = config.copy() # keep the original config untouched
assert isinstance( assert isinstance(

14
colossalai/builder/pipeline.py
View File

@ -163,17 +163,19 @@ def count_layer_params(layers):
def build_pipeline_model_from_cfg(config, num_chunks: int = 1, partition_method: str = 'parameter', verbose: bool = False): def build_pipeline_model_from_cfg(config, num_chunks: int = 1, partition_method: str = 'parameter', verbose: bool = False):
"""An intializer to split the model into different stages for pipeline parallelism. """An initializer to split the model into different stages for pipeline parallelism.
An example for the model config is shown below. The class VisionTransformerFromConfig should An example for the model config is shown below. The class VisionTransformerFromConfig should
inherit colossalai.nn.model.ModelFromConfig to allow this initializer to build model from a sequence inherit colossalai.nn.model.ModelFromConfig to allow this initializer to build model from a sequence
of layer configurations. of layer configurations.
model_config = dict( ::
type='VisionTransformerFromConfig',
embedding_cfg=dict(...), model_config = dict(
... type='VisionTransformerFromConfig',
) embedding_cfg=dict(...),
...
)
Args: Args:
config (dict): Configuration of the model. config (dict): Configuration of the model.

6
colossalai/communication/utils.py
View File

@ -45,7 +45,7 @@ def recv_tensor_meta(tensor_shape, prev_rank=None):
prev_rank (int): The rank of the source of the tensor. prev_rank (int): The rank of the source of the tensor.
Returns: Returns:
torch.Size: The shape of the tensor to be received. :class:`torch.Size`: The shape of the tensor to be received.
""" """
if tensor_shape is None: if tensor_shape is None:
if prev_rank is None: if prev_rank is None:
@ -71,7 +71,7 @@ def split_tensor_into_1d_equal_chunks(tensor, new_buffer=False):
new_buffer (bool, optional): Whether to use a new buffer to store sliced tensor. new_buffer (bool, optional): Whether to use a new buffer to store sliced tensor.
Returns: Returns:
torch.Tensor: The split tensor :class:`torch.Size`: The split tensor
""" """
partition_size = torch.numel(tensor) // gpc.get_world_size(ParallelMode.PARALLEL_1D) partition_size = torch.numel(tensor) // gpc.get_world_size(ParallelMode.PARALLEL_1D)
start_index = partition_size * gpc.get_local_rank(ParallelMode.PARALLEL_1D) start_index = partition_size * gpc.get_local_rank(ParallelMode.PARALLEL_1D)
@ -92,7 +92,7 @@ def gather_split_1d_tensor(tensor):
Args: Args:
tensor (torch.Tensor): Tensor to be gathered after communication. tensor (torch.Tensor): Tensor to be gathered after communication.
Returns: Returns:
gathered (torch.Tensor): The gathered tensor :class:`torch.Size`: The gathered tensor.
""" """
world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D) world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D)
numel = torch.numel(tensor) numel = torch.numel(tensor)

4
colossalai/context/parallel_context.py
View File

@ -193,7 +193,7 @@ class ParallelContext(metaclass=SingletonMeta):
Returns: Returns:
bool: a boolean value indicating whether the current device is the first one bool: a boolean value indicating whether the current device is the first one
among its group for `parallel_mode`. among its group for `parallel_mode`.
""" """
rank = self.get_local_rank(parallel_mode) rank = self.get_local_rank(parallel_mode)
return rank == 0 return rank == 0
@ -211,7 +211,7 @@ class ParallelContext(metaclass=SingletonMeta):
Returns: Returns:
bool: a boolean value indicating whether the current device is the first one bool: a boolean value indicating whether the current device is the first one
among its group for `parallel_mode`. among its group for `parallel_mode`.
""" """
rank = self.get_local_rank(parallel_mode) rank = self.get_local_rank(parallel_mode)
world_size = self.get_world_size(parallel_mode) world_size = self.get_world_size(parallel_mode)

7
colossalai/context/random/seed_manager.py
View File

@ -34,6 +34,7 @@ class SeedManager:
def set_state(self, parallel_mode: ParallelMode, state: Tensor): def set_state(self, parallel_mode: ParallelMode, state: Tensor):
"""Sets the state of the seed manager for `parallel_mode`. """Sets the state of the seed manager for `parallel_mode`.
Args: Args:
parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode. parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode.
state (:class:`torch.Tensor`): the state to be set. state (:class:`torch.Tensor`): the state to be set.
@ -66,9 +67,9 @@ class SeedManager:
seed (int): The seed to be added. seed (int): The seed to be added.
overwrtie (bool, optional): Whether allows to overwrite the seed that has been set already overwrtie (bool, optional): Whether allows to overwrite the seed that has been set already
Raises Raises:
AssertionError: Raises an AssertionError if `parallel_mode` is not an instance of AssertionError: Raises an AssertionError if `parallel_mode` is not an instance of :class:`colossalai.context.ParallelMode`
:class:`colossalai.context.ParallelMode` or the seed for `parallel_mode` has been added. or the seed for `parallel_mode` has been added.
""" """
assert isinstance(parallel_mode, ParallelMode), 'A valid ParallelMode must be provided' assert isinstance(parallel_mode, ParallelMode), 'A valid ParallelMode must be provided'
if overwrtie is False: if overwrtie is False:

2
colossalai/nn/layer/parallel_3d/_operation.py
View File

@ -264,7 +264,7 @@ def layernorm_3d(input_: Tensor, weight: Tensor, bias: Tensor, normalized_shape:
def split_tensor_3d(tensor: Tensor, dim: int, parallel_mode: ParallelMode) -> Tensor: def split_tensor_3d(tensor: Tensor, dim: int, parallel_mode: ParallelMode) -> Tensor:
r"""Splits 3D parallel tensor in specified dimension. r"""Splits 3D parallel tensor in specified dimension.
Args: Args:
tensor (:class:`torch.tensor`): Input tensor. tensor (:class:`torch.tensor`): Input tensor.
dim (int): Specified dimension in which to split. dim (int): Specified dimension in which to split.
parallel_mode (:class:`colossalai.context.parallel_mode.ParallelMode`, optional): Parallel mode. parallel_mode (:class:`colossalai.context.parallel_mode.ParallelMode`, optional): Parallel mode.

2
colossalai/nn/loss/loss_2d.py
View File

@ -27,7 +27,7 @@ class CrossEntropyLoss2D(_Loss):
reduce (bool, optional) reduce (bool, optional)
label_smoothing (float, optional) label_smoothing (float, optional)
More details about args, kwargs and torch.nn.functional.cross_entropy could be found in More details about ``args``, ``kwargs`` and ``torch.nn.functional.cross_entropy`` could be found in
`Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_. `Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_.
""" """

2
colossalai/nn/loss/loss_2p5d.py
View File

@ -27,7 +27,7 @@ class CrossEntropyLoss2p5D(_Loss):
reduce (bool, optional) reduce (bool, optional)
label_smoothing (float, optional) label_smoothing (float, optional)
More details about args, kwargs and torch.nn.functional.cross_entropy could be found in More details about ``args``, ``kwargs`` and ``torch.nn.functional.cross_entropy`` could be found in
`Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_. `Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_.
""" """
def __init__(self, reduction=True, *args, **kwargs): def __init__(self, reduction=True, *args, **kwargs):

2
colossalai/nn/loss/loss_3d.py
View File

@ -27,7 +27,7 @@ class CrossEntropyLoss3D(_Loss):
reduce (bool, optional) reduce (bool, optional)
label_smoothing (float, optional) label_smoothing (float, optional)
More details about args, kwargs and torch.nn.functional.cross_entropy could be found in More details about ``args``, ``kwargs`` and ``torch.nn.functional.cross_entropy`` could be found in
`Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_. `Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_.
""" """

4
colossalai/nn/loss/loss_moe.py
View File

@ -23,7 +23,7 @@ class MoeCrossEntropyLoss(_Loss):
reduction (str, optional) reduction (str, optional)
label_smoothing (float, optional) label_smoothing (float, optional)
More details about args, kwargs and torch.nn.functional.cross_entropy could be found in More details about ``args``, ``kwargs`` and ``torch.nn.functional.cross_entropy`` could be found in
`Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_. `Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_.
""" """
@ -40,7 +40,7 @@ class MoeCrossEntropyLoss(_Loss):
input (:class:`torch.tensor`): Predicted unnormalized scores (often referred to as logits). input (:class:`torch.tensor`): Predicted unnormalized scores (often referred to as logits).
target (:class:`torch.tensor`): Ground truth class indices or class probabilities. target (:class:`torch.tensor`): Ground truth class indices or class probabilities.
More details about args, kwargs and torch.nn.functional.cross_entropy could be found in More details about ``args``, ``kwargs`` and ``torch.nn.functional.cross_entropy`` could be found in
`Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_. `Cross_entropy <https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html#torch.nn.functional.cross_entropy>`_.
""" """
main_loss = self.loss(*args) main_loss = self.loss(*args)

3
colossalai/trainer/_trainer.py
View File

@ -307,8 +307,7 @@ class Trainer:
max_steps (int, optional): Maximum number of running iterations. max_steps (int, optional): Maximum number of running iterations.
test_dataloader (:class:`torch.utils.data.DataLoader`, optional): DataLoader for validation. test_dataloader (:class:`torch.utils.data.DataLoader`, optional): DataLoader for validation.
test_interval (int, optional): Interval of validation test_interval (int, optional): Interval of validation
hooks (list[`BaseHook <https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/trainer/hooks>`_], hooks (list[BaseHook], optional): A list of hooks used in training.
optional): A list of hooks used in training.
display_progress (bool, optional): If True, a progress bar will be displayed. display_progress (bool, optional): If True, a progress bar will be displayed.
""" """

27
colossalai/utils/memory_tracer/async_memtracer.py
View File

@ -21,21 +21,22 @@ class AsyncMemoryMonitor:
:type power: int :type power: int
Usage: Usage:
::
```python ```python
async_mem_monitor = AsyncMemoryMonitor() async_mem_monitor = AsyncMemoryMonitor()
input = torch.randn(2, 20).cuda() input = torch.randn(2, 20).cuda()
OP1 = torch.nn.Linear(20, 30).cuda() OP1 = torch.nn.Linear(20, 30).cuda()
OP2 = torch.nn.Linear(30, 40).cuda() OP2 = torch.nn.Linear(30, 40).cuda()
async_mem_monitor.start() async_mem_monitor.start()
output = OP1(input) output = OP1(input)
async_mem_monitor.finish() async_mem_monitor.finish()
async_mem_monitor.start() async_mem_monitor.start()
output = OP2(output) output = OP2(output)
async_mem_monitor.finish() async_mem_monitor.finish()
async_mem_monitor.save('log.pkl') async_mem_monitor.save('log.pkl')
``` ```
""" """
def __init__(self, power: int = 10): def __init__(self, power: int = 10):

31
colossalai/utils/profiler/prof_utils.py
View File

@ -73,25 +73,26 @@ class ProfilerContext(object):
""" """
Profiler context manager Profiler context manager
Usage: Usage:
::
```python ```python
world_size = 4 world_size = 4
inputs = torch.randn(10, 10, dtype=torch.float32, device=get_current_device()) inputs = torch.randn(10, 10, dtype=torch.float32, device=get_current_device())
outputs = torch.empty(world_size, 10, 10, dtype=torch.float32, device=get_current_device()) outputs = torch.empty(world_size, 10, 10, dtype=torch.float32, device=get_current_device())
outputs_list = list(torch.chunk(outputs, chunks=world_size, dim=0)) outputs_list = list(torch.chunk(outputs, chunks=world_size, dim=0))
cc_prof = CommProfiler() cc_prof = CommProfiler()
with ProfilerContext([cc_prof]) as prof: with ProfilerContext([cc_prof]) as prof:
op = dist.all_reduce(inputs, async_op=True) op = dist.all_reduce(inputs, async_op=True)
dist.all_gather(outputs_list, inputs) dist.all_gather(outputs_list, inputs)
op.wait() op.wait()
dist.reduce_scatter(inputs, outputs_list) dist.reduce_scatter(inputs, outputs_list)
dist.broadcast(inputs, 0) dist.broadcast(inputs, 0)
dist.reduce(inputs, 0) dist.reduce(inputs, 0)
prof.show() prof.show()
``` ```
""" """
def __init__(self, profilers: List[BaseProfiler] = None, enable: bool = True): def __init__(self, profilers: List[BaseProfiler] = None, enable: bool = True):

5
docs/colossalai/colossalai.amp.apex_amp.apex_amp.rst
View File

@ -1,5 +0,0 @@
colossalai.amp.apex\_amp.apex\_amp
==================================
.. automodule:: colossalai.amp.apex_amp.apex_amp
:members:

6
docs/colossalai/colossalai.amp.apex_amp.rst
View File

@ -3,9 +3,3 @@ colossalai.amp.apex\_amp
.. automodule:: colossalai.amp.apex_amp .. automodule:: colossalai.amp.apex_amp
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.amp.apex_amp.apex_amp

5
docs/colossalai/colossalai.amp.naive_amp.grad_scaler.base_grad_scaler.rst
View File

@ -1,5 +0,0 @@
colossalai.amp.naive\_amp.grad\_scaler.base\_grad\_scaler
=========================================================
.. automodule:: colossalai.amp.naive_amp.grad_scaler.base_grad_scaler
:members:

5
docs/colossalai/colossalai.amp.naive_amp.grad_scaler.constant_grad_scaler.rst
View File

@ -1,5 +0,0 @@
colossalai.amp.naive\_amp.grad\_scaler.constant\_grad\_scaler
=============================================================
.. automodule:: colossalai.amp.naive_amp.grad_scaler.constant_grad_scaler
:members:

5
docs/colossalai/colossalai.amp.naive_amp.grad_scaler.dynamic_grad_scaler.rst
View File

@ -1,5 +0,0 @@
colossalai.amp.naive\_amp.grad\_scaler.dynamic\_grad\_scaler
============================================================
.. automodule:: colossalai.amp.naive_amp.grad_scaler.dynamic_grad_scaler
:members:

13
docs/colossalai/colossalai.amp.naive_amp.grad_scaler.rst
View File

@ -1,13 +0,0 @@
colossalai.amp.naive\_amp.grad\_scaler
======================================
.. automodule:: colossalai.amp.naive_amp.grad_scaler
:members:
.. toctree::
:maxdepth: 2
colossalai.amp.naive_amp.grad_scaler.base_grad_scaler
colossalai.amp.naive_amp.grad_scaler.constant_grad_scaler
colossalai.amp.naive_amp.grad_scaler.dynamic_grad_scaler

5
docs/colossalai/colossalai.amp.naive_amp.naive_amp.rst
View File

@ -1,5 +0,0 @@
colossalai.amp.naive\_amp.naive\_amp
====================================
.. automodule:: colossalai.amp.naive_amp.naive_amp
:members:

11
docs/colossalai/colossalai.amp.naive_amp.rst
View File

@ -3,14 +3,3 @@ colossalai.amp.naive\_amp
.. automodule:: colossalai.amp.naive_amp .. automodule:: colossalai.amp.naive_amp
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.amp.naive_amp.grad_scaler
.. toctree::
:maxdepth: 2
colossalai.amp.naive_amp.naive_amp

6
docs/colossalai/colossalai.amp.rst
View File

@ -1,8 +1,7 @@
colossalai.amp colossalai.amp
============== ==============
.. automodule:: colossalai.amp *This module contains different amp models and optimizers.*
:members:
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
@ -16,3 +15,6 @@ colossalai.amp
:maxdepth: 2 :maxdepth: 2
colossalai.amp.amp_type colossalai.amp.amp_type
.. automodule:: colossalai.amp
:members: convert_to_amp

8
docs/colossalai/colossalai.amp.torch_amp.rst
View File

@ -2,10 +2,4 @@ colossalai.amp.torch\_amp
========================= =========================
.. automodule:: colossalai.amp.torch_amp .. automodule:: colossalai.amp.torch_amp
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.amp.torch_amp.torch_amp

5
docs/colossalai/colossalai.amp.torch_amp.torch_amp.rst
View File

@ -1,5 +0,0 @@
colossalai.amp.torch\_amp.torch\_amp
====================================
.. automodule:: colossalai.amp.torch_amp.torch_amp
:members:

5
docs/colossalai/colossalai.builder.rst
View File

@ -1,9 +1,8 @@
colossalai.builder colossalai.builder
================== ==================
.. automodule:: colossalai.builder *This module contains different methods to build* ``schedule``, ``model``, ``optimizer``,
:members: ``layers`` *and other important classes.*
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2

4
docs/colossalai/colossalai.communication.rst
View File

@ -1,9 +1,7 @@
colossalai.communication colossalai.communication
======================== ========================
.. automodule:: colossalai.communication *This module contains different methods for communication in parallel training.*
:members:
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2

5
docs/colossalai/colossalai.constants.rst
View File

@ -1,5 +0,0 @@
colossalai.constants
====================
.. automodule:: colossalai.constants
:members:

4
docs/colossalai/colossalai.context.moe_context.rst
View File

@ -1,5 +1,5 @@
colossalai.context.moe\_context colossalai.context.moe\_context
=============================== ====================================
.. automodule:: colossalai.context.moe_context .. automodule:: colossalai.context.moe_context
:members: :members:

2
docs/colossalai/colossalai.context.process_group_initializer.initializer_2d.rst
View File

@ -2,4 +2,4 @@ colossalai.context.process\_group\_initializer.initializer\_2d
============================================================== ==============================================================
.. automodule:: colossalai.context.process_group_initializer.initializer_2d .. automodule:: colossalai.context.process_group_initializer.initializer_2d
:members: :members: Initializer_2D_Row, Initializer_2D_Col

2
docs/colossalai/colossalai.context.process_group_initializer.initializer_2p5d.rst
View File

@ -2,4 +2,4 @@ colossalai.context.process\_group\_initializer.initializer\_2p5d
================================================================ ================================================================
.. automodule:: colossalai.context.process_group_initializer.initializer_2p5d .. automodule:: colossalai.context.process_group_initializer.initializer_2p5d
:members: :members: Initializer_2p5D_ROW, Initializer_2p5D_Col, Initializer_2p5D_Dep, Initializer_2p5D_XZ

2
docs/colossalai/colossalai.context.process_group_initializer.initializer_3d.rst
View File

@ -2,4 +2,4 @@ colossalai.context.process\_group\_initializer.initializer\_3d
============================================================== ==============================================================
.. automodule:: colossalai.context.process_group_initializer.initializer_3d .. automodule:: colossalai.context.process_group_initializer.initializer_3d
:members: :members: Initializer_3D_Input, Initializer_3D_Weight, Initializer_3D_Output

2
docs/colossalai/colossalai.context.process_group_initializer.initializer_sequence.rst
View File

@ -2,4 +2,4 @@ colossalai.context.process\_group\_initializer.initializer\_sequence
==================================================================== ====================================================================
.. automodule:: colossalai.context.process_group_initializer.initializer_sequence .. automodule:: colossalai.context.process_group_initializer.initializer_sequence
:members: :members: Initializer_Sequence_DP

1
docs/colossalai/colossalai.context.process_group_initializer.rst
View File

@ -14,6 +14,7 @@ colossalai.context.process\_group\_initializer
colossalai.context.process_group_initializer.initializer_3d colossalai.context.process_group_initializer.initializer_3d
colossalai.context.process_group_initializer.initializer_data colossalai.context.process_group_initializer.initializer_data
colossalai.context.process_group_initializer.initializer_model colossalai.context.process_group_initializer.initializer_model
colossalai.context.process_group_initializer.initializer_moe
colossalai.context.process_group_initializer.initializer_pipeline colossalai.context.process_group_initializer.initializer_pipeline
colossalai.context.process_group_initializer.initializer_sequence colossalai.context.process_group_initializer.initializer_sequence
colossalai.context.process_group_initializer.initializer_tensor colossalai.context.process_group_initializer.initializer_tensor

10
docs/colossalai/colossalai.context.rst
View File

@ -1,8 +1,8 @@
colossalai.context colossalai.context
================== ==================
.. automodule:: colossalai.context *This module is serving for receiving and passing user's configuration to all devices to
:members: initialize and construct parallel training.*
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
@ -15,7 +15,9 @@ colossalai.context
:maxdepth: 2 :maxdepth: 2
colossalai.context.config colossalai.context.config
colossalai.context.moe_context
colossalai.context.parallel_context colossalai.context.parallel_context
colossalai.context.moe_context
colossalai.context.parallel_mode colossalai.context.parallel_mode
colossalai.context.singleton_meta
.. automodule:: colossalai.context
:members:

5
docs/colossalai/colossalai.context.singleton_meta.rst
View File

@ -1,5 +0,0 @@
colossalai.context.singleton\_meta
==================================
.. automodule:: colossalai.context.singleton_meta
:members:

5
docs/colossalai/colossalai.core.rst
View File

@ -1,5 +0,0 @@
colossalai.core
===============
.. automodule:: colossalai.core
:members:

6
docs/colossalai/colossalai.engine.gradient_handler.rst
View File

@ -3,9 +3,3 @@ colossalai.engine.gradient\_handler
.. automodule:: colossalai.engine.gradient_handler .. automodule:: colossalai.engine.gradient_handler
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.engine.gradient_handler.utils

5
docs/colossalai/colossalai.engine.gradient_handler.utils.rst
View File

@ -1,5 +0,0 @@
colossalai.engine.gradient\_handler.utils
=========================================
.. automodule:: colossalai.engine.gradient_handler.utils
:members:

8
docs/colossalai/colossalai.engine.ophooks.rst
View File

@ -1,11 +1,5 @@
colossalai.engine.ophooks colossalai.engine.ophooks
========================= ===========================
.. automodule:: colossalai.engine.ophooks .. automodule:: colossalai.engine.ophooks
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.engine.ophooks.zero_hook

5
docs/colossalai/colossalai.engine.ophooks.zero_hook.rst
View File

@ -1,5 +0,0 @@
colossalai.engine.ophooks.zero\_hook
====================================
.. automodule:: colossalai.engine.ophooks.zero_hook
:members:

2
docs/colossalai/colossalai.engine.paramhooks.rst
View File

@ -1,5 +1,5 @@
colossalai.engine.paramhooks colossalai.engine.paramhooks
============================ ==============================
.. automodule:: colossalai.engine.paramhooks .. automodule:: colossalai.engine.paramhooks
:members: :members:

2
docs/colossalai/colossalai.engine.rst
View File

@ -8,6 +8,6 @@ colossalai.engine
:maxdepth: 2 :maxdepth: 2
colossalai.engine.gradient_handler colossalai.engine.gradient_handler
colossalai.engine.schedule
colossalai.engine.ophooks colossalai.engine.ophooks
colossalai.engine.paramhooks colossalai.engine.paramhooks
colossalai.engine.schedule

5
docs/colossalai/colossalai.global_variables.rst
View File

@ -1,5 +0,0 @@
colossalai.global\_variables
============================
.. automodule:: colossalai.global_variables
:members:

5
docs/colossalai/colossalai.logging.logger.rst
View File

@ -1,5 +0,0 @@
colossalai.logging.logger
=========================
.. automodule:: colossalai.logging.logger
:members:

6
docs/colossalai/colossalai.logging.rst
View File

@ -3,9 +3,3 @@ colossalai.logging
.. automodule:: colossalai.logging .. automodule:: colossalai.logging
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.logging.logger

5
docs/colossalai/colossalai.nn.layer.colossalai_layer.dropout.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.colossalai\_layer.dropout
=============================================
.. automodule:: colossalai.nn.layer.colossalai_layer.dropout
:members:

5
docs/colossalai/colossalai.nn.layer.colossalai_layer.embedding.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.colossalai\_layer.embedding
===============================================
.. automodule:: colossalai.nn.layer.colossalai_layer.embedding
:members:

5
docs/colossalai/colossalai.nn.layer.colossalai_layer.linear.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.colossalai\_layer.linear
============================================
.. automodule:: colossalai.nn.layer.colossalai_layer.linear
:members:

5
docs/colossalai/colossalai.nn.layer.colossalai_layer.normalization.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.colossalai\_layer.normalization
===================================================
.. automodule:: colossalai.nn.layer.colossalai_layer.normalization
:members:

9
docs/colossalai/colossalai.nn.layer.colossalai_layer.rst
View File

@ -3,12 +3,3 @@ colossalai.nn.layer.colossalai\_layer
.. automodule:: colossalai.nn.layer.colossalai_layer .. automodule:: colossalai.nn.layer.colossalai_layer
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.colossalai_layer.dropout
colossalai.nn.layer.colossalai_layer.embedding
colossalai.nn.layer.colossalai_layer.linear
colossalai.nn.layer.colossalai_layer.normalization

5
docs/colossalai/colossalai.nn.layer.moe.experts.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.moe.experts
===============================
.. automodule:: colossalai.nn.layer.moe.experts
:members:

5
docs/colossalai/colossalai.nn.layer.moe.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.moe.layers
==============================
.. automodule:: colossalai.nn.layer.moe.layers
:members:

8
docs/colossalai/colossalai.nn.layer.moe.rst
View File

@ -3,11 +3,3 @@ colossalai.nn.layer.moe
.. automodule:: colossalai.nn.layer.moe .. automodule:: colossalai.nn.layer.moe
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.moe.experts
colossalai.nn.layer.moe.layers
colossalai.nn.layer.moe.utils

5
docs/colossalai/colossalai.nn.layer.moe.utils.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.moe.utils
=============================
.. automodule:: colossalai.nn.layer.moe.utils
:members:

5
docs/colossalai/colossalai.nn.layer.parallel_1d.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.parallel\_1d.layers
=======================================
.. automodule:: colossalai.nn.layer.parallel_1d.layers
:members:

6
docs/colossalai/colossalai.nn.layer.parallel_1d.rst
View File

@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_1d
.. automodule:: colossalai.nn.layer.parallel_1d .. automodule:: colossalai.nn.layer.parallel_1d
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.parallel_1d.layers

5
docs/colossalai/colossalai.nn.layer.parallel_2d.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.parallel\_2d.layers
=======================================
.. automodule:: colossalai.nn.layer.parallel_2d.layers
:members:

8
docs/colossalai/colossalai.nn.layer.parallel_2d.rst
View File

@ -2,10 +2,4 @@ colossalai.nn.layer.parallel\_2d
================================ ================================
.. automodule:: colossalai.nn.layer.parallel_2d .. automodule:: colossalai.nn.layer.parallel_2d
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.parallel_2d.layers

5
docs/colossalai/colossalai.nn.layer.parallel_2p5d.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.parallel\_2p5d.layers
=========================================
.. automodule:: colossalai.nn.layer.parallel_2p5d.layers
:members:

6
docs/colossalai/colossalai.nn.layer.parallel_2p5d.rst
View File

@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_2p5d
.. automodule:: colossalai.nn.layer.parallel_2p5d .. automodule:: colossalai.nn.layer.parallel_2p5d
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.parallel_2p5d.layers

5
docs/colossalai/colossalai.nn.layer.parallel_3d.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.parallel\_3d.layers
=======================================
.. automodule:: colossalai.nn.layer.parallel_3d.layers
:members:

6
docs/colossalai/colossalai.nn.layer.parallel_3d.rst
View File

@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_3d
.. automodule:: colossalai.nn.layer.parallel_3d .. automodule:: colossalai.nn.layer.parallel_3d
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.parallel_3d.layers

5
docs/colossalai/colossalai.nn.layer.parallel_sequence.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.parallel\_sequence.layers
=============================================
.. automodule:: colossalai.nn.layer.parallel_sequence.layers
:members:

6
docs/colossalai/colossalai.nn.layer.parallel_sequence.rst
View File

@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_sequence
.. automodule:: colossalai.nn.layer.parallel_sequence .. automodule:: colossalai.nn.layer.parallel_sequence
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.parallel_sequence.layers

5
docs/colossalai/colossalai.nn.layer.utils.common.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.utils.common
================================
.. automodule:: colossalai.nn.layer.utils.common
:members:

6
docs/colossalai/colossalai.nn.layer.utils.rst
View File

@ -3,9 +3,3 @@ colossalai.nn.layer.utils
.. automodule:: colossalai.nn.layer.utils .. automodule:: colossalai.nn.layer.utils
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.utils.common

5
docs/colossalai/colossalai.nn.layer.vanilla.layers.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.vanilla.layers
==================================
.. automodule:: colossalai.nn.layer.vanilla.layers
:members:

6
docs/colossalai/colossalai.nn.layer.vanilla.rst
View File

@ -3,9 +3,3 @@ colossalai.nn.layer.vanilla
.. automodule:: colossalai.nn.layer.vanilla .. automodule:: colossalai.nn.layer.vanilla
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.vanilla.layers

5
docs/colossalai/colossalai.nn.layer.wrapper.lambda_wrapper.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.wrapper.lambda\_wrapper
===========================================
.. automodule:: colossalai.nn.layer.wrapper.lambda_wrapper
:members:

5
docs/colossalai/colossalai.nn.layer.wrapper.pipeline_wrapper.rst
View File

@ -1,5 +0,0 @@
colossalai.nn.layer.wrapper.pipeline\_wrapper
=============================================
.. automodule:: colossalai.nn.layer.wrapper.pipeline_wrapper
:members:

7
docs/colossalai/colossalai.nn.layer.wrapper.rst
View File

@ -3,10 +3,3 @@ colossalai.nn.layer.wrapper
.. automodule:: colossalai.nn.layer.wrapper .. automodule:: colossalai.nn.layer.wrapper
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.nn.layer.wrapper.lambda_wrapper
colossalai.nn.layer.wrapper.pipeline_wrapper

4
docs/colossalai/colossalai.nn.lr_scheduler.rst
View File

@ -1,10 +1,6 @@
colossalai.nn.lr\_scheduler colossalai.nn.lr\_scheduler
=========================== ===========================
.. automodule:: colossalai.nn.lr_scheduler
:members:
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2

4
docs/colossalai/colossalai.nn.optimizer.rst
View File

@ -1,10 +1,6 @@
colossalai.nn.optimizer colossalai.nn.optimizer
======================= =======================
.. automodule:: colossalai.nn.optimizer
:members:
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2

40
docs/colossalai/colossalai.nn.rst
View File

@ -1,21 +1,57 @@
colossalai.nn colossalai.nn
============= =============
.. automodule:: colossalai.nn *This part contains different colossalai layers for constructing your model.
:members: You can easily use them as the way of using layers in torch.nn.*
*Now colossalai support layer types below:* ``Linear``, ``Classifier``, ``Embedding``,
``PatchEmbedding``, ``LayerNorm``, ``Dropout`` *for different parallelisms.*
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
colossalai.nn.layer colossalai.nn.layer
*This part contains different loss functions for different parallelisms.*
.. toctree::
:maxdepth: 2
colossalai.nn.loss colossalai.nn.loss
*This part contains different learning rate schedules to control your learning rate
in training process*
.. toctree::
:maxdepth: 2
colossalai.nn.lr_scheduler colossalai.nn.lr_scheduler
*This part contains different metric to measure performance of your model.*
.. toctree::
:maxdepth: 2
colossalai.nn.metric colossalai.nn.metric
*This part contains some colossalai optimizers*
.. toctree::
:maxdepth: 2
colossalai.nn.model colossalai.nn.model
.. toctree::
:maxdepth: 2
colossalai.nn.optimizer colossalai.nn.optimizer
*This part contains different methods to initialize weights.*
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
colossalai.nn.init colossalai.nn.init
.. automodule:: colossalai.nn
:members:

9
docs/colossalai/colossalai.rst
View File

@ -1,9 +1,6 @@
colossalai Colossal-AI
========== ==========
.. automodule:: colossalai
:members:
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
@ -25,5 +22,7 @@ colossalai
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
colossalai.global_variables
colossalai.initialize colossalai.initialize
.. automodule:: colossalai
:members:

5
docs/colossalai/colossalai.testing.comparison.rst
View File

@ -1,5 +0,0 @@
colossalai.testing.comparison
=============================
.. automodule:: colossalai.testing.comparison
:members:

11
docs/colossalai/colossalai.testing.rst
View File

@ -1,12 +1,5 @@
colossalai.testing colossalai.testing
================== ========================
.. automodule:: colossalai.testing .. automodule:: colossalai.testing
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.testing.comparison
colossalai.testing.utils

5
docs/colossalai/colossalai.testing.utils.rst
View File

@ -1,5 +0,0 @@
colossalai.testing.utils
========================
.. automodule:: colossalai.testing.utils
:members:

5
docs/colossalai/colossalai.utils.data_sampler.base_sampler.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.data\_sampler.base\_sampler
============================================
.. automodule:: colossalai.utils.data_sampler.base_sampler
:members:

5
docs/colossalai/colossalai.utils.data_sampler.data_parallel_sampler.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.data\_sampler.data\_parallel\_sampler
======================================================
.. automodule:: colossalai.utils.data_sampler.data_parallel_sampler
:members:

7
docs/colossalai/colossalai.utils.data_sampler.rst
View File

@ -3,10 +3,3 @@ colossalai.utils.data\_sampler
.. automodule:: colossalai.utils.data_sampler .. automodule:: colossalai.utils.data_sampler
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.utils.data_sampler.base_sampler
colossalai.utils.data_sampler.data_parallel_sampler

5
docs/colossalai/colossalai.utils.memory_tracer.async_memtracer.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_tracer.async\_memtracer
================================================
.. automodule:: colossalai.utils.memory_tracer.async_memtracer
:members:

5
docs/colossalai/colossalai.utils.memory_tracer.memstats_collector.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_tracer.memstats\_collector
===================================================
.. automodule:: colossalai.utils.memory_tracer.memstats_collector
:members:

5
docs/colossalai/colossalai.utils.memory_tracer.model_data_memtracer.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_tracer.model\_data\_memtracer
======================================================
.. automodule:: colossalai.utils.memory_tracer.model_data_memtracer
:members:

14
docs/colossalai/colossalai.utils.memory_tracer.rst
View File

@ -1,15 +1,5 @@
colossalai.utils.memory\_tracer colossalai.utils.memory\_tracer
=============================== =======================
.. automodule:: colossalai.utils.memory_tracer .. automodule:: colossalai.utils.memory_tracer
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.utils.memory_tracer.async_memtracer
colossalai.utils.memory_tracer.memstats_collector
colossalai.utils.memory_tracer.model_data_memtracer
colossalai.utils.memory_tracer.test_async_memtracer
colossalai.utils.memory_tracer.test_memstats_collector

5
docs/colossalai/colossalai.utils.memory_tracer.test_async_memtracer.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_tracer.test\_async\_memtracer
======================================================
.. automodule:: colossalai.utils.memory_tracer.test_async_memtracer
:members:

5
docs/colossalai/colossalai.utils.memory_tracer.test_memstats_collector.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_tracer.test\_memstats\_collector
=========================================================
.. automodule:: colossalai.utils.memory_tracer.test_memstats_collector
:members:

5
docs/colossalai/colossalai.utils.memory_utils.bucket_tensor_copy.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_utils.bucket\_tensor\_copy
===================================================
.. automodule:: colossalai.utils.memory_utils.bucket_tensor_copy
:members:

5
docs/colossalai/colossalai.utils.memory_utils.memory_monitor.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_utils.memory\_monitor
==============================================
.. automodule:: colossalai.utils.memory_utils.memory_monitor
:members:

13
docs/colossalai/colossalai.utils.memory_utils.rst
View File

@ -1,13 +0,0 @@
colossalai.utils.memory\_utils
==============================
.. automodule:: colossalai.utils.memory_utils
:members:
.. toctree::
:maxdepth: 2
colossalai.utils.memory_utils.bucket_tensor_copy
colossalai.utils.memory_utils.memory_monitor
colossalai.utils.memory_utils.utils

5
docs/colossalai/colossalai.utils.memory_utils.utils.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.memory\_utils.utils
====================================
.. automodule:: colossalai.utils.memory_utils.utils
:members:

4
docs/colossalai/colossalai.utils.moe.rst
View File

@ -1,5 +1,5 @@
colossalai.utils.moe colossalai.utils.moe
==================== =======================
.. automodule:: colossalai.utils.moe .. automodule:: colossalai.utils.moe
:members: :members:

5
docs/colossalai/colossalai.utils.profiler.comm_profiler.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.profiler.comm\_profiler
========================================
.. automodule:: colossalai.utils.profiler.comm_profiler
:members:

5
docs/colossalai/colossalai.utils.profiler.pcie_profiler.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.profiler.pcie\_profiler
========================================
.. automodule:: colossalai.utils.profiler.pcie_profiler
:members:

5
docs/colossalai/colossalai.utils.profiler.prof_utils.rst
View File

@ -1,5 +0,0 @@
colossalai.utils.profiler.prof\_utils
=====================================
.. automodule:: colossalai.utils.profiler.prof_utils
:members:

12
docs/colossalai/colossalai.utils.profiler.rst
View File

@ -1,13 +1,5 @@
colossalai.utils.profiler colossalai.utils.profiler
========================= =======================
.. automodule:: colossalai.utils.profiler .. automodule:: colossalai.utils.profiler
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.utils.profiler.comm_profiler
colossalai.utils.profiler.pcie_profiler
colossalai.utils.profiler.prof_utils

7
docs/colossalai/colossalai.utils.rst
View File

@ -1,8 +1,7 @@
colossalai.utils colossalai.utils
================ ================
.. automodule:: colossalai.utils *This part contains useful utils for colossalai*
:members:
.. toctree:: .. toctree::
:maxdepth: 2 :maxdepth: 2
@ -10,9 +9,8 @@ colossalai.utils
colossalai.utils.data_sampler colossalai.utils.data_sampler
colossalai.utils.gradient_accumulation colossalai.utils.gradient_accumulation
colossalai.utils.memory_tracer colossalai.utils.memory_tracer
colossalai.utils.memory_utils
colossalai.utils.multi_tensor_apply colossalai.utils.multi_tensor_apply
colossalai.utils.profiler colossalai.utils.profile
colossalai.utils.tensor_detector colossalai.utils.tensor_detector
@ -23,5 +21,6 @@ colossalai.utils
colossalai.utils.checkpointing colossalai.utils.checkpointing
colossalai.utils.common colossalai.utils.common
colossalai.utils.cuda colossalai.utils.cuda
colossalai.utils.memory
colossalai.utils.moe colossalai.utils.moe
colossalai.utils.timer colossalai.utils.timer

12
docs/colossalai/colossalai.utils.tensor_detector.rst
View File

@ -1,11 +1,5 @@
colossalai.utils.tensor\_detector colossalai.utils.tensor\_detector
================================= =======================
.. automodule:: colossalai.utils.tensor_detector .. automodule:: colossalai.utils.tensor_detector.tensor_detector
:members: :members:
.. toctree::
:maxdepth: 2
colossalai.utils.tensor_detector.tensor_detector

Some files were not shown because too many files have changed in this diff Show More