diff --git a/docs/source/en/features/1D_tensor_parallel.md b/docs/source/en/features/1D_tensor_parallel.md index 79fe5ddea..0f01cfd32 100644 --- a/docs/source/en/features/1D_tensor_parallel.md +++ b/docs/source/en/features/1D_tensor_parallel.md @@ -2,14 +2,12 @@ Author: Zhengda Bian, Yongbin Li -> ⚠️ The information on this page is outdated and will be deprecated. Please check [Shardformer](./shardformer.md) for more information. - **Prerequisite** - [Define Your Configuration](../basics/define_your_config.md) - [Configure Parallelization](../basics/configure_parallelization.md) **Example Code** -- [ColossalAI-Examples 1D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) +- [Tensor Parallelism with Shardformer](https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/shardformer/examples) **Related Paper** - [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://deepakn94.github.io/assets/papers/megatron-sc21.pdf) @@ -44,79 +42,7 @@ Given $P$ processors, we present the theoretical computation and memory cost, as ## Usage -To enable 1D tensor parallelism for our model, e.g. on 2 GPUs, we need to configure the parallelism setting as below. -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=2, mode='1d'), -)) -``` -Then Colossal-AI will automatically apply 1D parallelism to all the layers from `colossalai.nn`. - -Let's define a model that consists of a two-layer multi-layer perceptron (MLP) as below. -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 - -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.transpose(0, 1).shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.transpose(0, 1).shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` - -Launch Colossal-AI on 2 GPUs and build the model. - -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -We will see the shapes of partitioned parameters(e.g. weights) in the MLP model. -```shell -Weight of the first linear layer: torch.Size([256, 512]) -Weight of the second linear layer: torch.Size([512, 256]) -``` -The complete weight of the first linear layer is supposed to have the shape `[256, 1024]`. After the column-parallel partitioning, it becomes `[256, 512]`. -Similarly, the second row-parallel layer partitions the weight `[1024, 256]` into `[512, 256]`. - -We can run the model with some random inputs. -```python -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -torch.distributed.broadcast(x, src=0) # synchronize input - -x = m(x) -``` -Then we can see the shapes of activation results. -```shell -Output of the first linear layer: torch.Size([16, 512]) -Output of the second linear layer: torch.Size([16, 256]) -``` -The output of the first linear layer is split into 2 partitions (each has the shape `[16, 512]`), while the second layer has identical outputs across the GPUs. +1D tensor parallelism is implemented by `Shardformer` feature in the newest version of ColossalAI. +For more details about ideas and usages of `Shardformer`, please refer to [Shardformer Doc](./shardformer.md). diff --git a/docs/source/en/features/2D_tensor_parallel.md b/docs/source/en/features/2D_tensor_parallel.md index aae8cc9ee..c79e7d196 100644 --- a/docs/source/en/features/2D_tensor_parallel.md +++ b/docs/source/en/features/2D_tensor_parallel.md @@ -60,83 +60,9 @@ Given $P=q\times q$ processors, we present the theoretical computation and memor ## Usage -To enable 2D tensor parallelism for our model, e.g. on 4 GPUs, we need to configure the parallelism setting as below. -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=4, mode='2d'), -)) -``` -Then Colossal-AI will automatically apply 2D parallelism to all the layers from `colossalai.nn`. +Currently the newest version of ColossalAI doesn't support 2D tensor parallelism, but this feature will be integrated into `Shardformer` in future releases. +For more details about ideas and usages of `Shardformer`, please refer to [Shardformer Doc](./shardformer.md). -Let's define a model that consists of a two-layer multi-layer perceptron (MLP) as below. -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 +For users of older version of ColossalAI, please refer to [ColossalAI-Examples - 2D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md). -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` -Launch Colossal-AI on 4 GPUs and build the model -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -We will see the shapes of partitioned parameters(e.g. weights) in the MLP model. -```shell -Weight of the first linear layer: torch.Size([128, 512]) -Weight of the second linear layer: torch.Size([512, 128]) -``` -The complete weight of the first linear layer is supposed to have the shape `[256, 1024]`. After the partitioning of 2D parallelism, it becomes `[128, 512]` on each GPU. -Similarly, the second layer partitions the weight `[1024, 256]` into `[512, 128]`. - -We can run the model with some random inputs. -```python -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -# partition input -torch.distributed.broadcast(x, src=0) -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)] -x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_ROW)] -print_rank_0(f'Input: {x.shape}') - -x = m(x) -``` -Then we can see the shapes of activation results. -```shell -Input: torch.Size([8, 128]) -Output of the first linear layer: torch.Size([8, 512]) -Output of the second linear layer: torch.Size([8, 128]) -``` -The activation tensors in 2D parallelism are all split in both row and column. -E.g. the output of the first linear layer has the shape `[8, 512]`, while the second layer has the output of `[8, 128]`. + diff --git a/docs/source/en/features/2p5D_tensor_parallel.md b/docs/source/en/features/2p5D_tensor_parallel.md index a81d14f10..b3cbd1c7c 100644 --- a/docs/source/en/features/2p5D_tensor_parallel.md +++ b/docs/source/en/features/2p5D_tensor_parallel.md @@ -58,86 +58,9 @@ Given $P=q \times q \times d$ processors, we present the theoretical computation ## Usage -To enable 2.5D tensor parallelism for our model, e.g. on 8 GPUs, we need to configure the parallelism setting as below. -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=8, mode='2.5d', depth=2), -)) +Currently the newest version of ColossalAI doesn't support 2.5D tensor parallelism, but this feature will be integrated into `Shardformer` in future releases. +For more details about ideas and usages of `Shardformer`, please refer to [Shardformer Doc](./shardformer.md). -``` -Then Colossal-AI will automatically apply 2.5D parallelism to all the layers from `colossalai.nn`. +For users of older version of ColossalAI, please refer to [ColossalAI-Examples - 2.5D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md). -Let's define a model that consists of a two-layer multi-layer perceptron (MLP) as below. -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 - -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` -Launch Colossal-AI on 8 GPUs and build the model -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -We will see the shapes of partitioned parameters(e.g. weights) in the MLP model. -```shell -Weight of the first linear layer: torch.Size([128, 512]) -Weight of the second linear layer: torch.Size([512, 128]) -``` -The complete weight of the first linear layer is supposed to have the shape `[256, 1024]`. After the partitioning of 2.5D parallelism, it becomes `[128, 512]` on each GPU. -Similarly, the second layer partitions the weight `[1024, 256]` into `[512, 128]`. - -We can run the model with some random inputs. -```python -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -# partition input -torch.distributed.broadcast(x, src=0) -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_DEP)] -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_COL)] -x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_ROW)] -print_rank_0(f'Input: {x.shape}') - -x = m(x) -``` -Then we can see the shapes of activation results. -```shell -Input: torch.Size([4, 128]) -Output of the first linear layer: torch.Size([4, 512]) -Output of the second linear layer: torch.Size([4, 128]) -``` -The activation tensors in 2.5D parallelism are all split by $d \times q$ in the row and $q$ in the column. -E.g. the output of the first linear layer has the shape `[4, 512]`), while the second layer has the output of `[4, 128]`. -Note, 2.5D parallelism use the same partition method as 2D parallelism for weights, where the difference is the partition of input. + diff --git a/docs/source/en/features/3D_tensor_parallel.md b/docs/source/en/features/3D_tensor_parallel.md index 0e28f08b2..00e6c5fca 100644 --- a/docs/source/en/features/3D_tensor_parallel.md +++ b/docs/source/en/features/3D_tensor_parallel.md @@ -67,85 +67,9 @@ Given $P=q \times q \times q$ processors, we present the theoretical computation ## Usage -To enable 3D tensor parallelism for our model, e.g. on 8 GPUs, we need to configure the parallelism setting as below. -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=8, mode='3d'), -)) -``` -Then Colossal-AI will automatically apply 3D parallelism to all the layers from `colossalai.nn`. +Currently the newest version of ColossalAI doesn't support 3D tensor parallelism, but this feature will be integrated into `Shardformer` in future releases. +For more details about ideas and usages of `Shardformer`, please refer to [Shardformer Doc](./shardformer.md). -Let's define a model that consists of a two-layer multi-layer perceptron (MLP) as below. -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 +For users of older version of ColossalAI, please refer to [ColossalAI-Examples - 3D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md). -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` -Launch Colossal-AI on 8 GPUs and build the model -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -We will see the shapes of partitioned parameters(e.g. weights) in the MLP model. -```shell -Weight of the first linear layer: torch.Size([128, 256]) -Weight of the second linear layer: torch.Size([512, 64]) -``` -The complete weight of the first linear layer is supposed to have the shape `[256, 1024]`. After the partitioning of 3D parallelism, it becomes `[128, 256]` on each GPU. -Similarly, the second layer partitions the weight `[1024, 256]` into `[512, 64]`. - -We can run the model with some random inputs. -```python -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -# partition input -torch.distributed.broadcast(x, src=0) -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_WEIGHT)] -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_INPUT)] -x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_OUTPUT)] -print_rank_0(f'Input: {x.shape}') - -x = m(x) -``` -Then we can see the shapes of activation results. -```shell -Input: torch.Size([4, 128]) -Output of the first linear layer: torch.Size([4, 512]) -Output of the second linear layer: torch.Size([4, 128]) -``` -The activation tensors in 3D parallelism are all split by $q^2$ in the row and $q$ in the column. -E.g. the output of the first linear layer has the shape `[4, 512]`), while the second layer has the output of `[4, 128]`. -Note, although the results of 3D parallelism have the same shape as that of 2.5D parallelism for weights here, the content of each partition is different. + diff --git a/docs/source/en/features/shardformer.md b/docs/source/en/features/shardformer.md index 872d00e4a..10e03e963 100644 --- a/docs/source/en/features/shardformer.md +++ b/docs/source/en/features/shardformer.md @@ -29,33 +29,6 @@ This module aims to make parallelization hassle-free for users who are not from Within a few lines of codes, users can turn a model into a state ready for distributed training. Also, Shardformer contains various optimization tools for acceleration and memory saving during forward/backward pass. - -## How Shardformer Works - -Generally, Shardformer works through the following four kinds of *replacements*: - -1. Replacing original PyTorch module (e.g. `nn.Linear`, `nn.Embedding`) with a crafted distributed module. -The distributed module keeps the same attributes as the original module but replaces the original parameters with distributed parameters. -Also, new `forward` methods will replace original ones so as to execute distributed computation, such as linear layers' split /gather operations executed under tensor parallelism. -Each distributed module implements its `from_native_module` static method to convert the PyTorch module to its corresponding distributed module. - -2. Replacing attributes of original Huggingface Transformers layers with appropriate attributes for distributed training. -For example, when training LlaMa-2 with tensor parallel size as 2, the attribute `num_heads` of `LlamaDecoderLayer` (the number of attention heads in each layer) should be replaced with `model.config.num_attention_heads // 2`. - -3. Replacing the `forward` methods implemented by original Huggingface -Transformers libraries with our customized `forward` methods. -This replacement is essential for pipeline paralellism, where a customiozed function is needed to pass intermediate hidden states between different pipeline stages. -Also, optimization methods such as flash attention or sequence parallel can be injected into the `forward` process through our customized `forward` method. - -4. Replacing the whole copy of model parameters and optimizer states with incomplete ones controlled by current device (this is why it's called Shardformer). -By executing `ModelSharder.shard` method, current device will only keep the part of model parameters it's supposed to take care of. -To be specific, they should be the assigned parameter shards when using tensor parallelism, or the parameters belonging to current pipeline stage when using pipeline parallelism, or both of them. -All other parameters are released so as to liberate memory usage. -As a result, the optimizer will only compute the states corresponding to these part of parameters, causing the usage of memory to be further saved. - -All of these replacements are implemented with manually written policies and forward functions. -If you want to delve deeper into the design of Shardformer or customize your own Shardformer policies, please refer to our [Shardformer development document](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/shardformer/README.md) and [pipeline parallelism design](https://github.com/hpcaitech/ColossalAI/discussions/4050) for more details. - ## Usage ### Shardformer Configuration @@ -101,31 +74,187 @@ is an example on how to trigger `Shardformer` through calling Shardformer APIs. ``` when training ChatGLM-2 with Shardformer, and initialize your model with these imported classes. +## How Shardformer Works + +Generally, Shardformer works through the following four kinds of *replacements*: + +1. Replacing original PyTorch module (e.g. `nn.Linear`, `nn.Embedding`) with a crafted distributed module. +The distributed module keeps the same attributes as the original module but replaces the original parameters with distributed parameters. +Also, new `forward` methods will replace original ones so as to execute distributed computation, such as linear layers' split /gather operations executed under tensor parallelism. +Each distributed module implements its `from_native_module` static method to convert the PyTorch module to its corresponding distributed module. + +2. Replacing attributes of original Huggingface Transformers layers with appropriate attributes for distributed training. +For example, when training LlaMa-2 with tensor parallel size as 2, the attribute `num_heads` of `LlamaDecoderLayer` (the number of attention heads in each layer) should be replaced with `model.config.num_attention_heads // 2`. + +3. Replacing the `forward` methods implemented by original Huggingface +Transformers libraries with our customized `forward` methods. +This replacement is essential for pipeline paralellism, where a customiozed function is needed to pass intermediate hidden states between different pipeline stages. +Also, optimization methods such as flash attention or sequence parallel can be injected into the `forward` process through our customized `forward` method. + +4. Replacing the whole copy of model parameters and optimizer states with incomplete ones controlled by current device (this is why it's called Shardformer). +By executing `ModelSharder.shard` method, current device will only keep the part of model parameters it's supposed to take care of. +To be specific, they should be the assigned parameter shards when using tensor parallelism, or the parameters belonging to current pipeline stage when using pipeline parallelism, or both of them. +All other parameters are released so as to liberate memory usage. +As a result, the optimizer will only compute the states corresponding to these part of parameters, causing the usage of memory to be further saved. + +All of these replacements are implemented with manually written policies and forward functions. +If you want to delve deeper into the design of Shardformer or customize your own Shardformer policies, please refer to our [Shardformer development document](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/shardformer/README.md) and [pipeline parallelism design](https://github.com/hpcaitech/ColossalAI/discussions/4050) for more details. ## Supporting Information -List of Huggingface transformers model families currently supported by Shardformer: -- LlaMa-1/LlaMa-2 -- GPT2 -- BERT -- OPT -- BLOOM -- T5 -- ViT -- ChatGLM-2 6B -- Whisper +Model/Feature Compatibility Matrix: -List of optimization tools currently supported by Shardformer: -- Flash Attention 2 -- JIT Fused Operator -- xFormers -- Fused Layer Normalization -- Sequence Parallel -- Sequence Overlap + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model/FeatureTensor
Parallel		Pipeline
Parallel		Lazy
Initialization		xFormersFlash
Attention 2		JIT Fused
Operators		Fused
LayerNorm		Sequence
Parallel		Sequence
Overlap
Llama V1/V2✔️✔️✔️✔️✔️✔️✔️❌❌
OPT✔️✔️✔️✔️✔️✔️✔️❌❌
BLOOM✔️✔️✔️✔️✔️✔️✔️✔️✔️
ChatGLM 2✔️✔️✔️✔️✔️✔️✔️✔️✔️
BERT✔️✔️✔️✔️✔️✔️✔️✔️✔️
GPT 2✔️✔️✔️✔️✔️✔️✔️✔️✔️
T5✔️✔️✔️✔️✔️✔️✔️❌❌
ViT✔️✔️❌✔️✔️✔️✔️❌❌
Whisper✔️✔️✔️✔️✔️❌✔️❌❌
SAM✔️❌❌✔️✔️✔️✔️❌❌
Blip2✔️❌❌✔️✔️✔️✔️❌❌
List of model families we plan to support in the near future: -- SAM -- Blip2 - RoBERTa - ALBERT - ERNIE @@ -135,9 +264,6 @@ List of model families we plan to support in the near future: - SwinTransformer V1/V2 - qwen -These lists will grow longer as more models and optimization tools emerge in the future. If you have any suggestions on the models/optimization we should support, please feel free to mention it in [Issues](https://github.com/hpcaitech/ColossalAI/issues) section of our project. - -For more details about compatibility between each optimization tool and each supported model, please refer to chapter Roadmap in our [develop document](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/shardformer/README.md). - +The support matrix will grow larger as more models and optimization tools emerge in the future. If you have any suggestions on the models/optimization we should support, please feel free to mention it in [Issues](https://github.com/hpcaitech/ColossalAI/issues) section of our project. diff --git a/docs/source/zh-Hans/features/1D_tensor_parallel.md b/docs/source/zh-Hans/features/1D_tensor_parallel.md index 61982cbb8..93fe9ea99 100644 --- a/docs/source/zh-Hans/features/1D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/1D_tensor_parallel.md @@ -2,14 +2,12 @@ 作者: Zhengda Bian, Yongbin Li -> ⚠️ 此页面上的信息已经过时并将被废弃。请在[Shardformer](./shardformer.md)页面查阅更新。 - **前置教程** - [定义配置文件](../basics/define_your_config.md) - [并行配置](../basics/configure_parallelization.md) -**示例代码** -- [ColossalAI-Examples 1D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md) +**示例代码**xw +- [Tensor Parallelism with Shardformer](https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/shardformer/examples) **相关论文** - [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://deepakn94.github.io/assets/papers/megatron-sc21.pdf) @@ -43,82 +41,10 @@ $$ | :-: | :-: | :-: | :-: | :-: | | $O(1/P)$ | $O(1/P)$ | $O(1)$ | $O(2(P-1)/P)$ | $O(2(P-1))$ | + ## 使用 -为了使模型能够实现一维张量并行, 如在2个 GPU 上, 我们需要配置如下的并行设置。 -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=2, mode='1d'), -)) -``` - -然后 Colossal-AI 会自动对所有来自 `colossalai.nn` 的层应用1D张量并行。 - -让我们定义一个由两层多层感知器 (MLP) 组成的模型，如下所示。 -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 - -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.transpose(0, 1).shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.transpose(0, 1).shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` - -在2个 GPU 上启动 Colossal-AI 并建立模型。 - -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -我们将会看到 MLP 模型中被划分的参数（如权重）的形状。 -```shell -Weight of the first linear layer: torch.Size([256, 512]) -Weight of the second linear layer: torch.Size([512, 256]) -``` -第一个线性层的完整权重形状应该为 `[256, 1024]`. 经过列-并行分割，它变成了 `[256, 512]`。 -同样地，第二个行并行层将权重 `[1024, 256]` 划分为 `[512, 256]`。 - -我们可以用一些随机输入来运行这个模型。 -```python -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -torch.distributed.broadcast(x, src=0) # synchronize input - -x = m(x) -``` -然后我们可以看到 activation 结果的形状。 -```shell -Output of the first linear layer: torch.Size([16, 512]) -Output of the second linear layer: torch.Size([16, 256]) -``` -第一个线性层的输出被划分成2块 (每个形状为 `[16, 512]`), 而第二层在整个 GPU 上的输出是相同的。 +在ColossalAI最新的版本中，1D张量并行由`Shardformer`功能实现。 +关于`Shardformer`的原理和用法细节请参考当前目录下的Shardformer文档。 diff --git a/docs/source/zh-Hans/features/2D_tensor_parallel.md b/docs/source/zh-Hans/features/2D_tensor_parallel.md index f163432ec..a8e5cf4bf 100644 --- a/docs/source/zh-Hans/features/2D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/2D_tensor_parallel.md @@ -60,82 +60,8 @@ $$ ## 使用 -为了使我们的模型能够实现二维张量并行，例如在4个 GPU 上，我们需要配置如下的并行设置。 -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=4, mode='2d'), -)) -``` -然后 Colossal-AI 会自动对所有来自 `colossalai.nn` 的层应用2D张量并行。 +ColossalAI的最新版本还暂不支持2D张量并行，但2D张量并行的功能会在未来的版本被集成入`Shardformer`中。关于`Shardformer`的原理和用法细节请参考当前目录下的Shardformer文档。 -让我们定义一个由两层多层感知器 (MLP) 组成的模型，如下所示。 -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 +对于老版本ColossalAI的用户，2D张量并行的用法请参考[ColossalAI-Examples - 2D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md)。 -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` -在4个 GPU 上启动 Colossal-AI 并建立模型。 -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -我们将会看到 MLP 模型中被划分的参数（如权重）的形状。 -```shell -Weight of the first linear layer: torch.Size([128, 512]) -Weight of the second linear layer: torch.Size([512, 128]) -``` -第一个线性层的完整权重形状应该为 `[256, 1024]`. 经过2D并行划分后，它在每个 GPU 上变成了 `[128, 512]` 。 -同样地，第二层将权重 `[1024, 256]` 划分为 `[512, 128]`. - -我们可以用一些随机输入来运行这个模型。 -```python -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -# partition input -torch.distributed.broadcast(x, src=0) -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_COL)] -x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_2D_ROW)] -print_rank_0(f'Input: {x.shape}') - -x = m(x) -``` -然后我们可以看到 activation 结果的形状。 -```shell -Input: torch.Size([8, 128]) -Output of the first linear layer: torch.Size([8, 512]) -Output of the second linear layer: torch.Size([8, 128]) -``` -2D并行中的 activation 张量都是同时在行和列分割的。例如，第一个线性层的输出是 `[8, 512]`, 而第二层的输出为 `[8, 128]`。 + diff --git a/docs/source/zh-Hans/features/2p5D_tensor_parallel.md b/docs/source/zh-Hans/features/2p5D_tensor_parallel.md index 5f1520272..6b0f1a301 100644 --- a/docs/source/zh-Hans/features/2p5D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/2p5D_tensor_parallel.md @@ -57,89 +57,8 @@ $$ ## 使用 -为了使我们的模型能够实现2.5D张量并行，例如在8个 GPU 上，我们需要配置如下的并行设置。 +ColossalAI的最新版本还暂不支持2.5D张量并行，但2.5D张量并行的功能会在未来的版本被集成入`Shardformer`中。关于`Shardformer`的原理和用法细节请参考当前目录下的Shardformer文档。 -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=8, mode='2.5d', depth=2), -)) +对于老版本ColossalAI的用户，2.5D张量并行的用法请参考[ColossalAI-Examples - 2.5D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md)。 -``` - -然后 Colossal-AI 会自动对所有来自 `colossalai.nn` 的层应用2.5D张量并行。 - -让我们定义一个由两层多层感知器 (MLP) 组成的模型，如下所示。 - -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 - -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` -在8个 GPU 上启动 Colossal-AI 并建立模型。 -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -我们将会看到 MLP 模型中被划分的参数（如权重）的形状。 -```shell -Weight of the first linear layer: torch.Size([128, 512]) -Weight of the second linear layer: torch.Size([512, 128]) -``` - -第一个线性层的完整权重形状应该为 `[256, 1024]`. 经过2.5D并行划分后，它在每个 GPU 上变成了 `[128, 512]` 。 -同样地，第二层将权重 `[1024, 256]` 划分为 `[512, 128]`. - -我们可以用一些随机输入来运行这个模型。 -```python -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -# partition input -torch.distributed.broadcast(x, src=0) -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_DEP)] -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_COL)] -x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_2P5D_ROW)] -print_rank_0(f'Input: {x.shape}') - -x = m(x) -``` -然后我们可以看到 activation 结果的形状。 -```shell -Input: torch.Size([4, 128]) -Output of the first linear layer: torch.Size([4, 512]) -Output of the second linear layer: torch.Size([4, 128]) -``` -2.5D并行中的 activation 张量都是同时在$d \times q$行和$q$列分割的。例如，第一个线性层的输出是 `[4, 512]`, 而第二层的输出为 `[4, 128]`。 -注意，2.5D并行使用与2D并行相同的划分方法来处理权重，区别在于对输入的划分。 + diff --git a/docs/source/zh-Hans/features/3D_tensor_parallel.md b/docs/source/zh-Hans/features/3D_tensor_parallel.md index 5ce0cdf6c..f6154559e 100644 --- a/docs/source/zh-Hans/features/3D_tensor_parallel.md +++ b/docs/source/zh-Hans/features/3D_tensor_parallel.md @@ -67,88 +67,8 @@ $$ ## 使用 -为了使我们的模型能够实现3D张量并行，例如在8个 GPU 上，我们需要配置如下的并行设置。 +ColossalAI的最新版本还暂不支持3D张量并行，但3D张量并行的功能会在未来的版本被集成入`Shardformer`中。关于`Shardformer`的原理和用法细节请参考当前目录下的Shardformer文档。 -```python -CONFIG = dict(parallel=dict( - data=1, - pipeline=1, - tensor=dict(size=8, mode='3d'), -)) -``` -然后 Colossal-AI 会自动对所有来自 `colossalai.nn` 的层应用3D张量并行。 +对于老版本ColossalAI的用户，3D张量并行的用法请参考[ColossalAI-Examples - 3D Tensor Parallelism](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/tensor_parallel/README.md)。 -让我们定义一个由两层多层感知器 (MLP) 组成的模型，如下所示。 - -```python -import colossalai -import colossalai.nn as col_nn -import torch -from colossalai.utils import print_rank_0 - -class MLP(torch.nn.Module): - def __init__(self, dim: int = 256): - super().__init__() - intermediate_dim = dim * 4 - self.dense_1 = col_nn.Linear(dim, intermediate_dim) - print_rank_0(f'Weight of the first linear layer: {self.dense_1.weight.shape}') - self.activation = torch.nn.GELU() - self.dense_2 = col_nn.Linear(intermediate_dim, dim) - print_rank_0(f'Weight of the second linear layer: {self.dense_2.weight.shape}') - self.dropout = col_nn.Dropout(0.1) - - def forward(self, x): - x = self.dense_1(x) - print_rank_0(f'Output of the first linear layer: {x.shape}') - x = self.activation(x) - x = self.dense_2(x) - print_rank_0(f'Output of the second linear layer: {x.shape}') - x = self.dropout(x) - return x -``` -在8个 GPU 上启动 Colossal-AI 并建立模型。 -```python -parser = colossalai.get_default_parser() -colossalai.launch(config=CONFIG, - rank=args.rank, - world_size=args.world_size, - local_rank=args.local_rank, - host=args.host, - port=args.port) - -m = MLP() -``` -我们将会看到 MLP 模型中被划分的参数（如权重）的形状。 -```shell -Weight of the first linear layer: torch.Size([128, 256]) -Weight of the second linear layer: torch.Size([512, 64]) -``` - -第一个线性层的完整权重形状应该为 `[256, 1024]`. 经过3D并行划分后，它在每个 GPU 上变成了 `[128, 256]` 。 -同样地，第二层将权重 `[1024, 256]` 划分为 `[512, 64]`. - -我们可以用一些随机输入来运行这个模型。 - -```python -from colossalai.context import ParallelMode -from colossalai.core import global_context as gpc -from colossalai.utils import get_current_device - -x = torch.randn((16, 256), device=get_current_device()) -# partition input -torch.distributed.broadcast(x, src=0) -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_WEIGHT)] -x = torch.chunk(x, 2, dim=0)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_INPUT)] -x = torch.chunk(x, 2, dim=-1)[gpc.get_local_rank(ParallelMode.PARALLEL_3D_OUTPUT)] -print_rank_0(f'Input: {x.shape}') - -x = m(x) -``` -然后我们可以看到 activation 结果的形状。 -```shell -Input: torch.Size([4, 128]) -Output of the first linear layer: torch.Size([4, 512]) -Output of the second linear layer: torch.Size([4, 128]) -``` -3D并行中的 activation 张量都是同时在$q^2$行和$q$列分割的。例如，第一个线性层的输出是 `[4, 512]`, 而第二层的输出为 `[4, 128]`。 -注意，虽然这里3D并行的结果与2.5D并行的结果形状相同，但每个划分的内容是不同的。 + diff --git a/docs/source/zh-Hans/features/shardformer.md b/docs/source/zh-Hans/features/shardformer.md index 49aa23e2d..e0d8df2c9 100644 --- a/docs/source/zh-Hans/features/shardformer.md +++ b/docs/source/zh-Hans/features/shardformer.md @@ -24,23 +24,6 @@ Author: [Baizhou Zhang](https://github.com/Fridge003) 出于这种动机，ColossalAI团队开发了**Shardformer**，该功能可以自动为HuggingFace中主流的Transformer模型进行封装，用于张量并行以及流水线并行的训练策略。如此一来，对系统了解不多的用户也可以轻松地在transformers模型上进行并行训练：只需几行代码，用户就可以将模型转变为并行训练的状态。此外，Shardformer也包括了多种优化工具，用于在前向/后向的传递过程中实现加速和节省内存。 - -## Shardformer的工作原理 - -通常来说，Shardformer通过以下四种“替换”进行工作： - -1. 用我们设计的分布式模块替换原始的PyTorch模块（例如`nn.Linear`、`nn.Embedding`）。 -分布式模块保持与原始模块相同的属性，但分布式模块会用新的参数替换原始模块的参数。新的前向函数将取代原来的前向函数，用于执行分布式计算，例如在张量并行下执行线性层的split/gather操作。每个分布式模块都应当实现其`from_native_module`静态方法，以将PyTorch模块转换为其相应的分布式模块。 - -2. 将原始Huggingface Transformers中间层的属性为适用于并行训练的属性。例如，当使用并行度为2的张量并行训练LlaMa-2时,`LlamaDecoderLayer` 的属性`num_heads`（每一层注意力头的数量）应替换为`model.config.num_attention_heads // 2`。 - -3. 将原来Huggingface transformers库实现的前向函数替换为我们定制的前向函数。前向函数的替换对于流水线并行性至关重要，因为流水线并行需要特殊的前向函数去在不同的流水线阶段之间传递中间的隐藏状态。此外，可以通过我们定制的前向函数将例如`flash attention`或序列并行的优化方法注入到前向的过程中。 - -4. 将完整的模型参数和优化器状态替换为只由当前设备控制的部分模型参数和优化器状态。通过执行`ModelSharder.shard`方法，当前设备仅会保留它应该处理的那部分模型参数。具体来说，这部分参数可以是使用张量并行时分配到当前机器的参数分片，或者使用流水线并行时当前流水线阶段的模型参数，或者兼而有之。除此之外的所有其他参数都被释放，用于节省内存的空间。 -如此一来，优化器只会计算保留的部分参数对应的状态，从而进一步节省内存的使用。 - -所有这些替换都是通过手动编写的策略和前向函数来实现的。如果您想更深入地研究Shardformer的设计方案，或者定制您自己的Shardformer策略，请参考[Shardformer 开发者文档](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/shardformer/README.md)和[流水并行设计方案](https://github.com/hpcaitech/ColossalAI/discussions/4050)以获得更多细节。 - ## 用法 ### Shardformer的参数配置 @@ -81,30 +64,179 @@ Shardformer的配置由类`ShardConfig`的参数控制： ``` 并且使用这些导入的类初始化模型。 + +## Shardformer的工作原理 + +通常来说，Shardformer通过以下四种“替换”进行工作： + +1. 用我们设计的分布式模块替换原始的PyTorch模块（例如`nn.Linear`、`nn.Embedding`）。 +分布式模块保持与原始模块相同的属性，但分布式模块会用新的参数替换原始模块的参数。新的前向函数将取代原来的前向函数，用于执行分布式计算，例如在张量并行下执行线性层的split/gather操作。每个分布式模块都应当实现其`from_native_module`静态方法，以将PyTorch模块转换为其相应的分布式模块。 + +2. 将原始Huggingface Transformers中间层的属性为适用于并行训练的属性。例如，当使用并行度为2的张量并行训练LlaMa-2时,`LlamaDecoderLayer` 的属性`num_heads`（每一层注意力头的数量）应替换为`model.config.num_attention_heads // 2`。 + +3. 将原来Huggingface transformers库实现的前向函数替换为我们定制的前向函数。前向函数的替换对于流水线并行性至关重要，因为流水线并行需要特殊的前向函数去在不同的流水线阶段之间传递中间的隐藏状态。此外，可以通过我们定制的前向函数将例如`flash attention`或序列并行的优化方法注入到前向的过程中。 + +4. 将完整的模型参数和优化器状态替换为只由当前设备控制的部分模型参数和优化器状态。通过执行`ModelSharder.shard`方法，当前设备仅会保留它应该处理的那部分模型参数。具体来说，这部分参数可以是使用张量并行时分配到当前机器的参数分片，或者使用流水线并行时当前流水线阶段的模型参数，或者兼而有之。除此之外的所有其他参数都被释放，用于节省内存的空间。 +如此一来，优化器只会计算保留的部分参数对应的状态，从而进一步节省内存的使用。 + +所有这些替换都是通过手动编写的策略和前向函数来实现的。如果您想更深入地研究Shardformer的设计方案，或者定制您自己的Shardformer策略，请参考[Shardformer 开发者文档](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/shardformer/README.md)和[流水并行设计方案](https://github.com/hpcaitech/ColossalAI/discussions/4050)以获得更多细节。 + + ## 支持信息 -Shardformer目前支持的Huggingface Transformer模型: -- LlaMa-1/LlaMa-2 -- GPT2 -- BERT -- OPT -- BLOOM -- T5 -- ViT -- ChatGLM-2 6B -- Whisper +模型/功能 兼容性矩阵： -Shardformer目前支持的优化工具: -- Flash Attention 2 -- JIT Fused Operator -- xFormers -- Fused Layer Normalization -- Sequence Parallel -- Sequence Overlap + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Model/FeatureTensor
Parallel		Pipeline
Parallel		Lazy
Initialization		xFormersFlash
Attention 2		JIT Fused
Operators		Fused
LayerNorm		Sequence
Parallel		Sequence
Overlap
Llama V1/V2✔️✔️✔️✔️✔️✔️✔️❌❌
OPT✔️✔️✔️✔️✔️✔️✔️❌❌
BLOOM✔️✔️✔️✔️✔️✔️✔️✔️✔️
ChatGLM 2✔️✔️✔️✔️✔️✔️✔️✔️✔️
BERT✔️✔️✔️✔️✔️✔️✔️✔️✔️
GPT 2✔️✔️✔️✔️✔️✔️✔️✔️✔️
T5✔️✔️✔️✔️✔️✔️✔️❌❌
ViT✔️✔️❌✔️✔️✔️✔️❌❌
Whisper✔️✔️✔️✔️✔️❌✔️❌❌
SAM✔️❌❌✔️✔️✔️✔️❌❌
Blip2✔️❌❌✔️✔️✔️✔️❌❌
我们计划在不久后为Shardformer支持的模型: -- SAM -- Blip2 - RoBERTa - ALBERT - ERNIE @@ -114,8 +246,6 @@ Shardformer目前支持的优化工具: - SwinTransformer V1/V2 - qwen -随着未来更多模型和优化工具的出现，这些列表将会变得越来越长。如果您对我们应该支持的模型/优化工具有任何建议，欢迎在项目的[Issues](https://github.com/hpcaitech/ColossalAI/issues)板块参与讨论。 - -更多关于不同优化工具和模型之间兼容性的细节，请参考[Shardformer开发者文档](https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/shardformer/README.md)中的Roadmap一节。 +随着未来更多模型和优化工具的出现，我们支持的模型/优化工具将会变得越来越多。如果您对我们应该支持的模型/优化工具有任何建议，欢迎在项目的[Issues](https://github.com/hpcaitech/ColossalAI/issues)板块参与讨论。