mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-17 23:18:36 +00:00
[doc] add shardformer support matrix/update tensor parallel documents (#4728)
* add compatibility matrix for shardformer doc * update tp doc
This commit is contained in:
Baizhou Zhang
GitHub
@@ -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文档。
|
||||
|
||||
<!-- doc-test-command: echo -->
|
||||
|
||||
@@ -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]`。
|
||||
<!-- doc-test-command: echo -->
|
||||
|
||||
@@ -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并行相同的划分方法来处理权重，区别在于对输入的划分。
|
||||
<!-- doc-test-command: echo -->
|
||||
|
||||
@@ -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并行的结果形状相同，但每个划分的内容是不同的。
|
||||
<!-- doc-test-command: echo -->
|
||||
|
||||
@@ -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
|
||||
<table>
|
||||
<tr>
|
||||
<th nowrap="nowrap">Model/Feature</th>
|
||||
<th nowrap="nowrap" title="Tensor Parallel">Tensor<br />Parallel</th>
|
||||
<th nowrap="nowrap" align="center" title="Pipeline Parallel">Pipeline<br />Parallel</th>
|
||||
<th nowrap="nowrap" align="center" title="Lazy Initialization">Lazy<br />Initialization</th>
|
||||
<th nowrap="nowrap" align="center" title="xFormers">xFormers</th>
|
||||
<th nowrap="nowrap" align="center" title="Flash Attention 2">Flash<br />Attention 2</th>
|
||||
<th nowrap="nowrap" align="center" title="JIT Fused Operators">JIT Fused<br />Operators</th>
|
||||
<th nowrap="nowrap" align="center" title="Fused LayerNorm">Fused<br />LayerNorm</th>
|
||||
<th nowrap="nowrap" align="center" title="Sequence Parallel">Sequence<br />Parallel</th>
|
||||
<th nowrap="nowrap" align="center" title="Sequence Overlap">Sequence<br />Overlap</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">Llama V1/V2</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">OPT</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">BLOOM</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">ChatGLM 2</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">BERT</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">GPT 2</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">T5</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">ViT</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">Whisper</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">SAM</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td nowrap="nowrap">Blip2</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">✔️</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
<td nowrap="nowrap" align="center">❌</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td colspan="39"></td>
|
||||
</tr>
|
||||
</table>
|
||||
|
||||
我们计划在不久后为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)板块参与讨论。
|
||||
|
||||
<!-- doc-test-command: echo -->
|
||||
|
||||
Reference in New Issue
Block a user