mirror of https://github.com/hpcaitech/ColossalAI
html refactor (#555)
Liang Bowen
GitHub
parent
d1211148a7
commit
2c45efc398
|
|
@ -19,7 +19,7 @@ def convert_to_amp(model: nn.Module, optimizer: Optimizer, criterion: _Loss, mod
|
|||
optimizer (:class:`torch.optim.Optimizer`): your optimizer object.
|
||||
criterion (:class:`torch.nn.modules.loss._Loss`): your loss function object.
|
||||
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:
|
||||
A tuple (model, optimizer, criterion).
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ def convert_to_apex_amp(model: nn.Module, optimizer: Optimizer, amp_config):
|
|||
Args:
|
||||
model (:class:`torch.nn.Module`): your model 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:
|
||||
::
|
||||
|
|
|
|||
|
|
@ -29,21 +29,22 @@ def build_from_registry(config, registry: Registry):
|
|||
is specified by `registry`.
|
||||
|
||||
Note:
|
||||
the `config` is used to construct the return object such as `LAYERS`,
|
||||
`OPTIMIZERS` and other support types in `registry`. The `config` should contain
|
||||
all required parameters of corresponding object. The details of support
|
||||
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>`_.
|
||||
the `config` is used to construct the return object such as `LAYERS`, `OPTIMIZERS`
|
||||
and other support types in `registry`. The `config` should contain
|
||||
all required parameters of corresponding object. The details of support
|
||||
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>`_.
|
||||
|
||||
Args:
|
||||
config (dict or :class:`colossalai.context.colossalai.context.Config`): information
|
||||
used in the construction 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:
|
||||
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
|
||||
assert isinstance(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""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
|
||||
inherit colossalai.nn.model.ModelFromConfig to allow this initializer to build model from a sequence
|
||||
of layer configurations.
|
||||
|
||||
model_config = dict(
|
||||
type='VisionTransformerFromConfig',
|
||||
embedding_cfg=dict(...),
|
||||
...
|
||||
)
|
||||
::
|
||||
|
||||
model_config = dict(
|
||||
type='VisionTransformerFromConfig',
|
||||
embedding_cfg=dict(...),
|
||||
...
|
||||
)
|
||||
|
||||
Args:
|
||||
config (dict): Configuration of the model.
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ def recv_tensor_meta(tensor_shape, prev_rank=None):
|
|||
prev_rank (int): The rank of the source of the tensor.
|
||||
|
||||
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 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.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The split tensor
|
||||
:class:`torch.Size`: The split tensor
|
||||
"""
|
||||
partition_size = torch.numel(tensor) // gpc.get_world_size(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:
|
||||
tensor (torch.Tensor): Tensor to be gathered after communication.
|
||||
Returns:
|
||||
gathered (torch.Tensor): The gathered tensor
|
||||
:class:`torch.Size`: The gathered tensor.
|
||||
"""
|
||||
world_size = gpc.get_world_size(ParallelMode.PARALLEL_1D)
|
||||
numel = torch.numel(tensor)
|
||||
|
|
|
|||
|
|
@ -193,7 +193,7 @@ class ParallelContext(metaclass=SingletonMeta):
|
|||
|
||||
Returns:
|
||||
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)
|
||||
return rank == 0
|
||||
|
|
@ -211,7 +211,7 @@ class ParallelContext(metaclass=SingletonMeta):
|
|||
|
||||
Returns:
|
||||
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)
|
||||
world_size = self.get_world_size(parallel_mode)
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ class SeedManager:
|
|||
|
||||
def set_state(self, parallel_mode: ParallelMode, state: Tensor):
|
||||
"""Sets the state of the seed manager for `parallel_mode`.
|
||||
|
||||
Args:
|
||||
parallel_mode (:class:`colossalai.context.ParallelMode`): The chosen parallel mode.
|
||||
state (:class:`torch.Tensor`): the state to be set.
|
||||
|
|
@ -66,9 +67,9 @@ class SeedManager:
|
|||
seed (int): The seed to be added.
|
||||
overwrtie (bool, optional): Whether allows to overwrite the seed that has been set already
|
||||
|
||||
Raises
|
||||
AssertionError: Raises an AssertionError if `parallel_mode` is not an instance of
|
||||
:class:`colossalai.context.ParallelMode` or the seed for `parallel_mode` has been added.
|
||||
Raises:
|
||||
AssertionError: Raises an AssertionError if `parallel_mode` is not an instance of :class:`colossalai.context.ParallelMode`
|
||||
or the seed for `parallel_mode` has been added.
|
||||
"""
|
||||
assert isinstance(parallel_mode, ParallelMode), 'A valid ParallelMode must be provided'
|
||||
if overwrtie is False:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
r"""Splits 3D parallel tensor in specified dimension.
|
||||
|
||||
Args:
|
||||
Args:
|
||||
tensor (:class:`torch.tensor`): Input tensor.
|
||||
dim (int): Specified dimension in which to split.
|
||||
parallel_mode (:class:`colossalai.context.parallel_mode.ParallelMode`, optional): Parallel mode.
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ class CrossEntropyLoss2D(_Loss):
|
|||
reduce (bool, 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>`_.
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ class CrossEntropyLoss2p5D(_Loss):
|
|||
reduce (bool, 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>`_.
|
||||
"""
|
||||
def __init__(self, reduction=True, *args, **kwargs):
|
||||
|
|
|
|||
|
|
@ -27,7 +27,7 @@ class CrossEntropyLoss3D(_Loss):
|
|||
reduce (bool, 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>`_.
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ class MoeCrossEntropyLoss(_Loss):
|
|||
reduction (str, 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>`_.
|
||||
"""
|
||||
|
||||
|
|
@ -40,7 +40,7 @@ class MoeCrossEntropyLoss(_Loss):
|
|||
input (:class:`torch.tensor`): Predicted unnormalized scores (often referred to as logits).
|
||||
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>`_.
|
||||
"""
|
||||
main_loss = self.loss(*args)
|
||||
|
|
|
|||
|
|
@ -307,8 +307,7 @@ class Trainer:
|
|||
max_steps (int, optional): Maximum number of running iterations.
|
||||
test_dataloader (:class:`torch.utils.data.DataLoader`, optional): DataLoader for validation.
|
||||
test_interval (int, optional): Interval of validation
|
||||
hooks (list[`BaseHook <https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/trainer/hooks>`_],
|
||||
optional): A list of hooks used in training.
|
||||
hooks (list[BaseHook], optional): A list of hooks used in training.
|
||||
display_progress (bool, optional): If True, a progress bar will be displayed.
|
||||
"""
|
||||
|
||||
|
|
|
|||
|
|
@ -21,21 +21,22 @@ class AsyncMemoryMonitor:
|
|||
:type power: int
|
||||
|
||||
Usage:
|
||||
::
|
||||
|
||||
```python
|
||||
async_mem_monitor = AsyncMemoryMonitor()
|
||||
input = torch.randn(2, 20).cuda()
|
||||
OP1 = torch.nn.Linear(20, 30).cuda()
|
||||
OP2 = torch.nn.Linear(30, 40).cuda()
|
||||
```python
|
||||
async_mem_monitor = AsyncMemoryMonitor()
|
||||
input = torch.randn(2, 20).cuda()
|
||||
OP1 = torch.nn.Linear(20, 30).cuda()
|
||||
OP2 = torch.nn.Linear(30, 40).cuda()
|
||||
|
||||
async_mem_monitor.start()
|
||||
output = OP1(input)
|
||||
async_mem_monitor.finish()
|
||||
async_mem_monitor.start()
|
||||
output = OP2(output)
|
||||
async_mem_monitor.finish()
|
||||
async_mem_monitor.save('log.pkl')
|
||||
```
|
||||
async_mem_monitor.start()
|
||||
output = OP1(input)
|
||||
async_mem_monitor.finish()
|
||||
async_mem_monitor.start()
|
||||
output = OP2(output)
|
||||
async_mem_monitor.finish()
|
||||
async_mem_monitor.save('log.pkl')
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, power: int = 10):
|
||||
|
|
|
|||
|
|
@ -73,25 +73,26 @@ class ProfilerContext(object):
|
|||
"""
|
||||
Profiler context manager
|
||||
Usage:
|
||||
::
|
||||
|
||||
```python
|
||||
world_size = 4
|
||||
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_list = list(torch.chunk(outputs, chunks=world_size, dim=0))
|
||||
```python
|
||||
world_size = 4
|
||||
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_list = list(torch.chunk(outputs, chunks=world_size, dim=0))
|
||||
|
||||
cc_prof = CommProfiler()
|
||||
cc_prof = CommProfiler()
|
||||
|
||||
with ProfilerContext([cc_prof]) as prof:
|
||||
op = dist.all_reduce(inputs, async_op=True)
|
||||
dist.all_gather(outputs_list, inputs)
|
||||
op.wait()
|
||||
dist.reduce_scatter(inputs, outputs_list)
|
||||
dist.broadcast(inputs, 0)
|
||||
dist.reduce(inputs, 0)
|
||||
with ProfilerContext([cc_prof]) as prof:
|
||||
op = dist.all_reduce(inputs, async_op=True)
|
||||
dist.all_gather(outputs_list, inputs)
|
||||
op.wait()
|
||||
dist.reduce_scatter(inputs, outputs_list)
|
||||
dist.broadcast(inputs, 0)
|
||||
dist.reduce(inputs, 0)
|
||||
|
||||
prof.show()
|
||||
```
|
||||
prof.show()
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, profilers: List[BaseProfiler] = None, enable: bool = True):
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.amp.apex\_amp.apex\_amp
|
||||
==================================
|
||||
|
||||
.. automodule:: colossalai.amp.apex_amp.apex_amp
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.amp.apex\_amp
|
|||
|
||||
.. automodule:: colossalai.amp.apex_amp
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.amp.apex_amp.apex_amp
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.amp.naive\_amp.grad\_scaler.base\_grad\_scaler
|
||||
=========================================================
|
||||
|
||||
.. automodule:: colossalai.amp.naive_amp.grad_scaler.base_grad_scaler
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.amp.naive\_amp.grad\_scaler.constant\_grad\_scaler
|
||||
=============================================================
|
||||
|
||||
.. automodule:: colossalai.amp.naive_amp.grad_scaler.constant_grad_scaler
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.amp.naive\_amp.grad\_scaler.dynamic\_grad\_scaler
|
||||
============================================================
|
||||
|
||||
.. automodule:: colossalai.amp.naive_amp.grad_scaler.dynamic_grad_scaler
|
||||
:members:
|
||||
|
|
@ -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
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.amp.naive\_amp.naive\_amp
|
||||
====================================
|
||||
|
||||
.. automodule:: colossalai.amp.naive_amp.naive_amp
|
||||
:members:
|
||||
|
|
@ -3,14 +3,3 @@ colossalai.amp.naive\_amp
|
|||
|
||||
.. automodule:: colossalai.amp.naive_amp
|
||||
:members:
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.amp.naive_amp.grad_scaler
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.amp.naive_amp.naive_amp
|
||||
|
|
|
|||
|
|
@ -1,8 +1,7 @@
|
|||
colossalai.amp
|
||||
==============
|
||||
|
||||
.. automodule:: colossalai.amp
|
||||
:members:
|
||||
*This module contains different amp models and optimizers.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
|
@ -16,3 +15,6 @@ colossalai.amp
|
|||
:maxdepth: 2
|
||||
|
||||
colossalai.amp.amp_type
|
||||
|
||||
.. automodule:: colossalai.amp
|
||||
:members: convert_to_amp
|
||||
|
|
|
|||
|
|
@ -2,10 +2,4 @@ colossalai.amp.torch\_amp
|
|||
=========================
|
||||
|
||||
.. automodule:: colossalai.amp.torch_amp
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.amp.torch_amp.torch_amp
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.amp.torch\_amp.torch\_amp
|
||||
====================================
|
||||
|
||||
.. automodule:: colossalai.amp.torch_amp.torch_amp
|
||||
:members:
|
||||
|
|
@ -1,9 +1,8 @@
|
|||
colossalai.builder
|
||||
==================
|
||||
|
||||
.. automodule:: colossalai.builder
|
||||
:members:
|
||||
|
||||
*This module contains different methods to build* ``schedule``, ``model``, ``optimizer``,
|
||||
``layers`` *and other important classes.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
colossalai.communication
|
||||
========================
|
||||
|
||||
.. automodule:: colossalai.communication
|
||||
:members:
|
||||
|
||||
*This module contains different methods for communication in parallel training.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.constants
|
||||
====================
|
||||
|
||||
.. automodule:: colossalai.constants
|
||||
:members:
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
colossalai.context.moe\_context
|
||||
===============================
|
||||
====================================
|
||||
|
||||
.. automodule:: colossalai.context.moe_context
|
||||
:members:
|
||||
:members:
|
||||
|
|
@ -2,4 +2,4 @@ colossalai.context.process\_group\_initializer.initializer\_2d
|
|||
==============================================================
|
||||
|
||||
.. automodule:: colossalai.context.process_group_initializer.initializer_2d
|
||||
:members:
|
||||
:members: Initializer_2D_Row, Initializer_2D_Col
|
||||
|
|
|
|||
|
|
@ -2,4 +2,4 @@ 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,4 +2,4 @@ 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,4 +2,4 @@ colossalai.context.process\_group\_initializer.initializer\_sequence
|
|||
====================================================================
|
||||
|
||||
.. automodule:: colossalai.context.process_group_initializer.initializer_sequence
|
||||
:members:
|
||||
:members: Initializer_Sequence_DP
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ colossalai.context.process\_group\_initializer
|
|||
colossalai.context.process_group_initializer.initializer_3d
|
||||
colossalai.context.process_group_initializer.initializer_data
|
||||
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_sequence
|
||||
colossalai.context.process_group_initializer.initializer_tensor
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
colossalai.context
|
||||
==================
|
||||
|
||||
.. automodule:: colossalai.context
|
||||
:members:
|
||||
*This module is serving for receiving and passing user's configuration to all devices to
|
||||
initialize and construct parallel training.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
|
@ -15,7 +15,9 @@ colossalai.context
|
|||
:maxdepth: 2
|
||||
|
||||
colossalai.context.config
|
||||
colossalai.context.moe_context
|
||||
colossalai.context.parallel_context
|
||||
colossalai.context.moe_context
|
||||
colossalai.context.parallel_mode
|
||||
colossalai.context.singleton_meta
|
||||
|
||||
.. automodule:: colossalai.context
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.context.singleton\_meta
|
||||
==================================
|
||||
|
||||
.. automodule:: colossalai.context.singleton_meta
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.core
|
||||
===============
|
||||
|
||||
.. automodule:: colossalai.core
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.engine.gradient\_handler
|
|||
|
||||
.. automodule:: colossalai.engine.gradient_handler
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.engine.gradient_handler.utils
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.engine.gradient\_handler.utils
|
||||
=========================================
|
||||
|
||||
.. automodule:: colossalai.engine.gradient_handler.utils
|
||||
:members:
|
||||
|
|
@ -1,11 +1,5 @@
|
|||
colossalai.engine.ophooks
|
||||
=========================
|
||||
===========================
|
||||
|
||||
.. automodule:: colossalai.engine.ophooks
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.engine.ophooks.zero_hook
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.engine.ophooks.zero\_hook
|
||||
====================================
|
||||
|
||||
.. automodule:: colossalai.engine.ophooks.zero_hook
|
||||
:members:
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
colossalai.engine.paramhooks
|
||||
============================
|
||||
==============================
|
||||
|
||||
.. automodule:: colossalai.engine.paramhooks
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -8,6 +8,6 @@ colossalai.engine
|
|||
:maxdepth: 2
|
||||
|
||||
colossalai.engine.gradient_handler
|
||||
colossalai.engine.schedule
|
||||
colossalai.engine.ophooks
|
||||
colossalai.engine.paramhooks
|
||||
colossalai.engine.schedule
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.global\_variables
|
||||
============================
|
||||
|
||||
.. automodule:: colossalai.global_variables
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.logging.logger
|
||||
=========================
|
||||
|
||||
.. automodule:: colossalai.logging.logger
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.logging
|
|||
|
||||
.. automodule:: colossalai.logging
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.logging.logger
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.colossalai\_layer.dropout
|
||||
=============================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.colossalai_layer.dropout
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.colossalai\_layer.embedding
|
||||
===============================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.colossalai_layer.embedding
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.colossalai\_layer.linear
|
||||
============================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.colossalai_layer.linear
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.colossalai\_layer.normalization
|
||||
===================================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.colossalai_layer.normalization
|
||||
:members:
|
||||
|
|
@ -3,12 +3,3 @@ colossalai.nn.layer.colossalai\_layer
|
|||
|
||||
.. automodule:: colossalai.nn.layer.colossalai_layer
|
||||
: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
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.moe.experts
|
||||
===============================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.moe.experts
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.moe.layers
|
||||
==============================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.moe.layers
|
||||
:members:
|
||||
|
|
@ -3,11 +3,3 @@ colossalai.nn.layer.moe
|
|||
|
||||
.. automodule:: colossalai.nn.layer.moe
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.moe.experts
|
||||
colossalai.nn.layer.moe.layers
|
||||
colossalai.nn.layer.moe.utils
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.moe.utils
|
||||
=============================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.moe.utils
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.parallel\_1d.layers
|
||||
=======================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.parallel_1d.layers
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_1d
|
|||
|
||||
.. automodule:: colossalai.nn.layer.parallel_1d
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.parallel_1d.layers
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.parallel\_2d.layers
|
||||
=======================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.parallel_2d.layers
|
||||
:members:
|
||||
|
|
@ -2,10 +2,4 @@ colossalai.nn.layer.parallel\_2d
|
|||
================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.parallel_2d
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.parallel_2d.layers
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.parallel\_2p5d.layers
|
||||
=========================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.parallel_2p5d.layers
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_2p5d
|
|||
|
||||
.. automodule:: colossalai.nn.layer.parallel_2p5d
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.parallel_2p5d.layers
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.parallel\_3d.layers
|
||||
=======================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.parallel_3d.layers
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_3d
|
|||
|
||||
.. automodule:: colossalai.nn.layer.parallel_3d
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.parallel_3d.layers
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.parallel\_sequence.layers
|
||||
=============================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.parallel_sequence.layers
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.nn.layer.parallel\_sequence
|
|||
|
||||
.. automodule:: colossalai.nn.layer.parallel_sequence
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.parallel_sequence.layers
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.utils.common
|
||||
================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.utils.common
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.nn.layer.utils
|
|||
|
||||
.. automodule:: colossalai.nn.layer.utils
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.utils.common
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.vanilla.layers
|
||||
==================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.vanilla.layers
|
||||
:members:
|
||||
|
|
@ -3,9 +3,3 @@ colossalai.nn.layer.vanilla
|
|||
|
||||
.. automodule:: colossalai.nn.layer.vanilla
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.vanilla.layers
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.wrapper.lambda\_wrapper
|
||||
===========================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.wrapper.lambda_wrapper
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.nn.layer.wrapper.pipeline\_wrapper
|
||||
=============================================
|
||||
|
||||
.. automodule:: colossalai.nn.layer.wrapper.pipeline_wrapper
|
||||
:members:
|
||||
|
|
@ -3,10 +3,3 @@ colossalai.nn.layer.wrapper
|
|||
|
||||
.. automodule:: colossalai.nn.layer.wrapper
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer.wrapper.lambda_wrapper
|
||||
colossalai.nn.layer.wrapper.pipeline_wrapper
|
||||
|
|
|
|||
|
|
@ -1,10 +1,6 @@
|
|||
colossalai.nn.lr\_scheduler
|
||||
===========================
|
||||
|
||||
.. automodule:: colossalai.nn.lr_scheduler
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
|
|
|
|||
|
|
@ -1,10 +1,6 @@
|
|||
colossalai.nn.optimizer
|
||||
=======================
|
||||
|
||||
.. automodule:: colossalai.nn.optimizer
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
|
|
|
|||
|
|
@ -1,21 +1,57 @@
|
|||
colossalai.nn
|
||||
=============
|
||||
|
||||
.. automodule:: colossalai.nn
|
||||
:members:
|
||||
*This part contains different colossalai layers for constructing your model.
|
||||
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::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.layer
|
||||
|
||||
*This part contains different loss functions for different parallelisms.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
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
|
||||
|
||||
*This part contains different metric to measure performance of your model.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.metric
|
||||
|
||||
*This part contains some colossalai optimizers*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.model
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.optimizer
|
||||
|
||||
*This part contains different methods to initialize weights.*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.nn.init
|
||||
|
||||
.. automodule:: colossalai.nn
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -1,9 +1,6 @@
|
|||
colossalai
|
||||
Colossal-AI
|
||||
==========
|
||||
|
||||
.. automodule:: colossalai
|
||||
:members:
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
|
|
@ -25,5 +22,7 @@ colossalai
|
|||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.global_variables
|
||||
colossalai.initialize
|
||||
|
||||
.. automodule:: colossalai
|
||||
:members:
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.testing.comparison
|
||||
=============================
|
||||
|
||||
.. automodule:: colossalai.testing.comparison
|
||||
:members:
|
||||
|
|
@ -1,12 +1,5 @@
|
|||
colossalai.testing
|
||||
==================
|
||||
========================
|
||||
|
||||
.. automodule:: colossalai.testing
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.testing.comparison
|
||||
colossalai.testing.utils
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.testing.utils
|
||||
========================
|
||||
|
||||
.. automodule:: colossalai.testing.utils
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.data\_sampler.base\_sampler
|
||||
============================================
|
||||
|
||||
.. automodule:: colossalai.utils.data_sampler.base_sampler
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.data\_sampler.data\_parallel\_sampler
|
||||
======================================================
|
||||
|
||||
.. automodule:: colossalai.utils.data_sampler.data_parallel_sampler
|
||||
:members:
|
||||
|
|
@ -3,10 +3,3 @@ colossalai.utils.data\_sampler
|
|||
|
||||
.. automodule:: colossalai.utils.data_sampler
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.utils.data_sampler.base_sampler
|
||||
colossalai.utils.data_sampler.data_parallel_sampler
|
||||
|
|
|
|||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_tracer.async\_memtracer
|
||||
================================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_tracer.async_memtracer
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_tracer.memstats\_collector
|
||||
===================================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_tracer.memstats_collector
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_tracer.model\_data\_memtracer
|
||||
======================================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_tracer.model_data_memtracer
|
||||
:members:
|
||||
|
|
@ -1,15 +1,5 @@
|
|||
colossalai.utils.memory\_tracer
|
||||
===============================
|
||||
=======================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_tracer
|
||||
: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
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_tracer.test\_async\_memtracer
|
||||
======================================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_tracer.test_async_memtracer
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_tracer.test\_memstats\_collector
|
||||
=========================================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_tracer.test_memstats_collector
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_utils.bucket\_tensor\_copy
|
||||
===================================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_utils.bucket_tensor_copy
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_utils.memory\_monitor
|
||||
==============================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_utils.memory_monitor
|
||||
:members:
|
||||
|
|
@ -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
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.memory\_utils.utils
|
||||
====================================
|
||||
|
||||
.. automodule:: colossalai.utils.memory_utils.utils
|
||||
:members:
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
colossalai.utils.moe
|
||||
====================
|
||||
=======================
|
||||
|
||||
.. automodule:: colossalai.utils.moe
|
||||
:members:
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.profiler.comm\_profiler
|
||||
========================================
|
||||
|
||||
.. automodule:: colossalai.utils.profiler.comm_profiler
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.profiler.pcie\_profiler
|
||||
========================================
|
||||
|
||||
.. automodule:: colossalai.utils.profiler.pcie_profiler
|
||||
:members:
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
colossalai.utils.profiler.prof\_utils
|
||||
=====================================
|
||||
|
||||
.. automodule:: colossalai.utils.profiler.prof_utils
|
||||
:members:
|
||||
|
|
@ -1,13 +1,5 @@
|
|||
colossalai.utils.profiler
|
||||
=========================
|
||||
=======================
|
||||
|
||||
.. automodule:: colossalai.utils.profiler
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.utils.profiler.comm_profiler
|
||||
colossalai.utils.profiler.pcie_profiler
|
||||
colossalai.utils.profiler.prof_utils
|
||||
:members:
|
||||
|
|
@ -1,8 +1,7 @@
|
|||
colossalai.utils
|
||||
================
|
||||
|
||||
.. automodule:: colossalai.utils
|
||||
:members:
|
||||
*This part contains useful utils for colossalai*
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
|
@ -10,9 +9,8 @@ colossalai.utils
|
|||
colossalai.utils.data_sampler
|
||||
colossalai.utils.gradient_accumulation
|
||||
colossalai.utils.memory_tracer
|
||||
colossalai.utils.memory_utils
|
||||
colossalai.utils.multi_tensor_apply
|
||||
colossalai.utils.profiler
|
||||
colossalai.utils.profile
|
||||
colossalai.utils.tensor_detector
|
||||
|
||||
|
||||
|
|
@ -23,5 +21,6 @@ colossalai.utils
|
|||
colossalai.utils.checkpointing
|
||||
colossalai.utils.common
|
||||
colossalai.utils.cuda
|
||||
colossalai.utils.memory
|
||||
colossalai.utils.moe
|
||||
colossalai.utils.timer
|
||||
|
|
|
|||
|
|
@ -1,11 +1,5 @@
|
|||
colossalai.utils.tensor\_detector
|
||||
=================================
|
||||
=======================
|
||||
|
||||
.. automodule:: colossalai.utils.tensor_detector
|
||||
:members:
|
||||
|
||||
|
||||
.. toctree::
|
||||
:maxdepth: 2
|
||||
|
||||
colossalai.utils.tensor_detector.tensor_detector
|
||||
.. automodule:: colossalai.utils.tensor_detector.tensor_detector
|
||||
:members:
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue