mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-03 01:55:12 +00:00
[doc] migrate the markdown files (#2652)
This commit is contained in:
97
docs/source/en/basics/colotensor_concept.md
Normal file
97
docs/source/en/basics/colotensor_concept.md
Normal file
@@ -0,0 +1,97 @@
|
||||
# ColoTensor Concepts
|
||||
|
||||
Author: [Jiarui Fang](https://github.com/feifeibear), [Hongxin Liu](https://github.com/ver217) and [Haichen Huang](https://github.com/1SAA)
|
||||
|
||||
**Prerequisite:**
|
||||
- [Colossal-AI Overview](../concepts/colossalai_overview.md)
|
||||
- [Distributed Training](../concepts/distributed_training.md)
|
||||
- [Paradigms of Parallelism](../concepts/paradigms_of_parallelism.md)
|
||||
|
||||
## Introduction
|
||||
|
||||
After ColossalAI version 0.1.8, [ColoTensor](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.html#colossalai.tensor.ColoTensor) becomes the basic data structure for tensors in ColossalAI. It is a subclass of torch.Tensor and can be used as a PyTorch Tensor. Additionally, some unique features make it possible to represent a Global Tensor with a payload distributed across multiple GPU devices. With the help of ColoTensor, the users can write distributed DNN training program similar to a serial one.support the following features.
|
||||
|
||||
ColoTensor contains extra attributes capsuled in a [ColoTensorSpec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.tensor_spec.html#colossalai.tensor.tensor_spec.ColoTensorSpec) instance to describe the tensor's payload distribution and computing pattern.
|
||||
|
||||
- ProcessGroup: how processes are organized as communication groups.
|
||||
- Distributed Spec: how tensor is distributed among process groups.
|
||||
- Compute Spec: how the tensor is used during computation.
|
||||
|
||||
We elaborate on them one by one.
|
||||
|
||||
## ProcessGroup
|
||||
|
||||
An instance of class [ProcessGroup](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.html#colossalai.tensor.ProcessGroup) describes how processes are organized in process groups. Processes in a process group can participate in the same collective communication operations together, such as allgather, allreduce, etc. The way the process group is organized is dominated by the Tensor's parallelism strategy. For example, if the user defines the tensor parallel (TP) and data parallel (DP) modes of a tensor, then the process organization of the process group will be automatically deduced. The process group settings can vary among different tensors. Therefore, it enables us to support more complicated hybrid parallel. The pipeline parallel (PP) definition is not in the ProcessGroup, it needs another set of mechanisms . We will supplement the related content of ColoTensor applied to PP in the future.
|
||||
|
||||
Currently, a process group of ColoTensor is defined by two configurations, i.e. tp_degree and dp_degree. In the case of DP+TP hybrid parallelism, the device can be viewed as a 2D mesh. We place TP communication groups on the leading low dimension of the device mesh and then place the data parallel groups along the high dimension of the device mesh. The reason is that tensor parallelism has a larger communication overhead than data parallelism. Neighboring devices are placed inside a TP process group and are often placed in the same node.
|
||||
|
||||
Considering that 8 processes are configured as tp_degree=4, and dp_degree=2, the layout is shown below. Process group tp0 contains gpu 0,1,2,3. Process dp1 contains gpu 1 and 5.
|
||||
|
||||
<figure style={{textAlign: "center"}}>
|
||||
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/colossalai/img/ColoTensor_layout_demo.PNG"/>
|
||||
<figcaption>Process Group using tp_degree=4, dp_degree=2</figcaption>
|
||||
</figure>
|
||||
|
||||
## Distributed Spec
|
||||
|
||||
An instance of [Distributed Spec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.distspec.html) describes how a ColoTensor is distributed among the ProcessGroup.
|
||||
|
||||
How tensors are distributed among DP process groups is automatically derived and does not need to be manually specified by the user. If this tensor is a model parameter, it is replicated within the DP process group. If it is an activation tensor, it is split along the process with the highest dimension and evenly distributed the tensor payload among processes in the DP process group.
|
||||
|
||||
Therefore, when using Distributed Spec, we only need to describe the way that the tensor is distributed among TP process groups. There are currently two ways to distribute among TP process group, i.e. [ShardSpec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.distspec.html#colossalai.tensor.distspec.ShardSpec) and [ReplicaSpec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.distspec.html#colossalai.tensor.distspec.ReplicaSpec). ShardSpec needs to specify the dimension index dim of the partition and the number of partitions num_partitions. Currently, we only support the split on a single dim. Different dist specs on the TP process groups can be converted to each other through the set_dist_spec() interface. The spec conversions are recorded by the autograd mechanism and it will trigger corresponding reverse operations during backward propagation.
|
||||
|
||||
## Compute Spec
|
||||
|
||||
An instance of class [ComputeSpec](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.compute_spec.html#colossalai.tensor.compute_spec.ComputeSpec) describes how a Coloensor be used in DNN training. Currently, we will set the correct Compute Pattern for the ColoTensor as the parameters of the module. The specific application scenarios will be shown in the next document.
|
||||
|
||||
## ColoParameter
|
||||
|
||||
[ColoParameter](https://colossalai.readthedocs.io/en/latest/colossalai/colossalai.tensor.colo_parameter.html#colossalai.tensor.colo_parameter.ColoParameter) is a subclass of ColoTensor. Used to define a Global Parameter tensor. Its relationship with ColoTensor is consistent with Torch.Tensor and torch.Parameter. The latter allows the tensor to appear in the return values of the module's parameters() and name_parameters() methods.
|
||||
|
||||
## Example
|
||||
|
||||
Let's see an example. A ColoTensor is initialized and sharded on 8 GPUs using tp_degree=4, dp_dgree=2. And then the tensor is sharded along the last dim among the TP process groups. Finally, we reshard it along the first dim (0 dim) among the TP process groups. We encourage users to run the code and observe the shape of each tensor.
|
||||
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch.multiprocessing as mp
|
||||
from colossalai.utils import free_port, print_rank_0
|
||||
from functools import partial
|
||||
|
||||
import colossalai
|
||||
from colossalai.tensor import ProcessGroup, ColoTensor, ColoTensorSpec, ShardSpec, ComputeSpec, ComputePattern
|
||||
from colossalai.utils import free_port
|
||||
|
||||
import torch
|
||||
|
||||
def run_dist_tests(rank, world_size, port):
|
||||
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
|
||||
pg = ProcessGroup(tp_degree=2, dp_degree=2)
|
||||
|
||||
torch.manual_seed(0)
|
||||
local_tensor = torch.randn(2, 3, 1).cuda()
|
||||
print_rank_0(f"shape {local_tensor.shape}, {local_tensor.data}")
|
||||
|
||||
spec = ColoTensorSpec(pg, ShardSpec(dims=[-1], num_partitions=[pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D))
|
||||
t1 = ColoTensor.from_torch_tensor(local_tensor, spec)
|
||||
t1 = t1.to_replicate()
|
||||
print_rank_0(f"shape {t1.shape}, {t1.data}")
|
||||
|
||||
spec2 = ShardSpec([0], [pg.tp_world_size()])
|
||||
t1.set_dist_spec(spec2)
|
||||
print_rank_0(f"shape {t1.shape}, {t1.data}")
|
||||
|
||||
def test_dist_cases(world_size):
|
||||
run_func = partial(run_dist_tests, world_size=world_size, port=free_port())
|
||||
mp.spawn(run_func, nprocs=world_size)
|
||||
|
||||
if __name__ == '__main__':
|
||||
test_dist_cases(4)
|
||||
```
|
||||
|
||||
:::caution
|
||||
|
||||
The ColoTensor is an experimental feature and may be updated.
|
||||
|
||||
:::
|
53
docs/source/en/basics/command_line_tool.md
Normal file
53
docs/source/en/basics/command_line_tool.md
Normal file
@@ -0,0 +1,53 @@
|
||||
# Command Line Tool
|
||||
|
||||
Author: Shenggui Li
|
||||
|
||||
**Prerequisite:**
|
||||
- [Distributed Training](../concepts/distributed_training.md)
|
||||
- [Colossal-AI Overview](../concepts/colossalai_overview.md)
|
||||
|
||||
## Introduction
|
||||
|
||||
Colossal-AI provides command-line utilities for the user.
|
||||
The current command line tools support the following features.
|
||||
|
||||
- verify Colossal-AI build
|
||||
- launch distributed jobs
|
||||
- tensor parallel micro-benchmarking
|
||||
|
||||
## Check Installation
|
||||
|
||||
To verify whether your Colossal-AI is built correctly, you can use the command `colossalai check -i`.
|
||||
This command will inform you information regarding the version compatibility and cuda extension.
|
||||
|
||||
<figure style={{textAlign: "center"}}>
|
||||
<img src="https://s2.loli.net/2022/05/04/KJmcVknyPHpBofa.png"/>
|
||||
<figcaption>Check Installation Demo</figcaption>
|
||||
</figure>
|
||||
|
||||
## Launcher
|
||||
|
||||
To launch distributed jobs on single or multiple nodes, the command `colossalai run` can be used for process launching.
|
||||
You may refer to [Launch Colossal-AI](./launch_colossalai.md) for more details.
|
||||
|
||||
## Tensor Parallel Micro-Benchmarking
|
||||
|
||||
As Colossal-AI provides an array of tensor parallelism methods, it is not intuitive to choose one for your hardware and
|
||||
model. Therefore, we provide a simple benchmarking to evaluate the performance of various tensor parallelisms on your system.
|
||||
This benchmarking is run on a simple MLP model where the input data is of the shape `(batch_size, seq_length, hidden_size)`.
|
||||
Based on the number of GPUs, the CLI will look for all possible tensor parallel configurations and display the benchmarking results.
|
||||
You can customize the benchmarking configurations by checking out `colossalai benchmark --help`.
|
||||
|
||||
```shell
|
||||
# run on 4 GPUs
|
||||
colossalai benchmark --gpus 4
|
||||
|
||||
# run on 8 GPUs
|
||||
colossalai benchmark --gpus 8
|
||||
```
|
||||
|
||||
:::caution
|
||||
|
||||
Only single-node benchmarking is supported currently.
|
||||
|
||||
:::
|
156
docs/source/en/basics/configure_parallelization.md
Normal file
156
docs/source/en/basics/configure_parallelization.md
Normal file
@@ -0,0 +1,156 @@
|
||||
# Configure Parallelization
|
||||
|
||||
Author: Shenggui Li, Siqi Mai
|
||||
|
||||
**Prerequisite:**
|
||||
- [Distributed Training](../concepts/distributed_training.md)
|
||||
- [Paradigms of Parallelism](../concepts/paradigms_of_parallelism.md)
|
||||
- [Define Your Configuration](./define_your_config.md)
|
||||
|
||||
|
||||
## Introduction
|
||||
|
||||
We support multiple parallelization in Colossal-AI. Hybrid parallelism in our codebase refers to namely the combination
|
||||
of data parallelism, pipeline parallelism and tensor parallelism (1D, 2D, 2.5D, 3D).
|
||||
|
||||
Each parallelism requires different network topology and thus initialize different process groups.
|
||||
You can initialize the corresponding process group by setting `parallel` in the config file.
|
||||
The configuration for `parallel` must obey the following format. Data parallel size will be
|
||||
inferred automatically based on your inputs to pipeline parallelism and tensor parallelism.
|
||||
`colossalai.launch` will initialize these distributed process groups automatically based on your configuration.
|
||||
|
||||
Some sample configurations are shown below:
|
||||
|
||||
```python
|
||||
# sampler format
|
||||
parallel = dict(
|
||||
pipeline=dict("size": int),
|
||||
tensor=dict("size": int, "mode": '1d' or '2d' or '2.5d' or '3d', "kwargs": Any)
|
||||
)
|
||||
|
||||
# this is ok
|
||||
parallel = dict(
|
||||
pipeline=dict(size=2),
|
||||
tensor=dict(size=4, mode='2d')
|
||||
)
|
||||
|
||||
# this is ok
|
||||
parallel = dict(
|
||||
pipeline=2,
|
||||
tensor=dict(size=4, mode='2d')
|
||||
)
|
||||
|
||||
# this is not ok
|
||||
# as you need to specify the mode for tensor parallelism
|
||||
parallel = dict(
|
||||
pipeline=2,
|
||||
tensor=4
|
||||
)
|
||||
|
||||
# this is ok as well as tensor will be default to size 1
|
||||
# and mode None
|
||||
parallel = dict(
|
||||
pipeline=2
|
||||
)
|
||||
|
||||
# this is ok as well as pipeline will default to size 1
|
||||
parallel = dict(
|
||||
tensor=dict(size=4, mode='2d')
|
||||
)
|
||||
|
||||
```
|
||||
|
||||
The key name `size` refers to the parallel size of the parallelism dimension. For example, pipeline size 2 means there
|
||||
will be 2 pipeline stages. The key name `mode` in tensor parallel config means the corresponding tensor parallelism
|
||||
will be initialized.
|
||||
|
||||
**You can choose to not have 'parallel' in your configuration and both pipeline and tensor will default to size 1.**
|
||||
|
||||
**Total number of GPUs must be equal to `data parallel size * tensor parallel size * pipeline parallel size`**
|
||||
|
||||
## Data Parallel
|
||||
|
||||
Data parallel is the most common way to distribute your training task by splitting data into several shards and train on
|
||||
a single shard on each device. The configuration for data parallel is detected automatically and set for you. You do not
|
||||
have to explicitly set them in your configurations. There are two ways to handle the all-reduce in data parallel in Colossal-AI.
|
||||
|
||||
1. If you specify gradient handlers, gradients will be all-reduced according to the gradient handlers
|
||||
2. Otherwise, PyTorch DistributedDataParallel will be used
|
||||
|
||||
In most cases, you will be using the second mode unless you have complex handling of the gradients.
|
||||
|
||||
## 1D, 2D, 2.5D and 3D Parallel
|
||||
|
||||
To enable hybrid parallelism, we provide an array of tensor parallelism. We provide the list of papers which match each
|
||||
tensor parallel method. These parallel modes need to work with the distributed layers provided by Colossal-AI.
|
||||
|
||||
- 1D: [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://arxiv.org/abs/1909.08053)
|
||||
|
||||
- 2D: [An Efficient 2D Method for Training Super-Large Deep Learning Models](https://arxiv.org/abs/2104.05343)
|
||||
2D parallel relies on the SUMMA matrix multiplication algorithm and splits the input data, model weights and layer
|
||||
outputs along two different dimensions. The tensor chunks are distributed over a 2D mesh of `P = N^2` devices where
|
||||
`N` is the number of tensor chunks in a single dimension.
|
||||
|
||||
- 2.5D: [2.5-dimensional distributed model training](https://arxiv.org/abs/2105.14500)
|
||||
Inspired by the 2.5D matrix multiplication algorithm, 2.5D parallel introduces a novel tensor parallelism which
|
||||
further parallelizes 2D tensor parallelism. An amount of `P = N^2 ∗ d` processors are arranged into `d` layers, where
|
||||
each layer performs matrix multiplication operations independently with a dimension `N`.
|
||||
|
||||
- 3D: [Maximizing Parallelism in Distributed Training for Huge Neural Networks](https://arxiv.org/abs/2105.14450)
|
||||
We also introduce a 3D tensor parallelism that parallelizes neural networks on a 3D processor cube. This method
|
||||
achieves the optimal, `O(P^{1/3})` communication overhead on $P$ processors, while both computation and memory usage
|
||||
are evenly distributed through optimized load balancing of parameters as well as activations.
|
||||
|
||||
```python
|
||||
# 1D parallel
|
||||
parallel = dict(
|
||||
tensor=dict(size=4, mode='1d')
|
||||
)
|
||||
|
||||
# 2D parallel
|
||||
parallel = dict(
|
||||
tensor=dict(size=4, mode='2d')
|
||||
)
|
||||
|
||||
# 2.5D parallel
|
||||
parallel = dict(
|
||||
tensor=dict(size=8, mode='2.5d', depth=2)
|
||||
)
|
||||
|
||||
# 3D parallel
|
||||
parallel = dict(
|
||||
tensor=dict(size=8, mode='3d')
|
||||
)
|
||||
```
|
||||
|
||||
Once you specify the tensor parallel mode in your configuration, you can proceed to use its corresponding distributed
|
||||
operator. For example, if you mode is '2d', you can use `colossalai.nn.Linear2D` in you model construction.
|
||||
|
||||
|
||||
## Pipeline Parallel
|
||||
|
||||
Pipeline parallelism is to split the model into several partitions by layer. For example, let's assume we have a simple
|
||||
model which consists of two linear layer. We have two GPUs, and we can allocate the first linear layer to the first GPU
|
||||
and the second layer to the second GPU.
|
||||
|
||||
You can set the number of pipeline stages in your configuration file. When pipeline size is larger than 1, Colossal-AI
|
||||
will automatically creates the pipeline schedule which defines the forward and backward step.
|
||||
|
||||
```python
|
||||
parallel = dict(
|
||||
pipeline=dict(size=4), # number of pipeline stages
|
||||
)
|
||||
```
|
||||
|
||||
## Sequence Parallel
|
||||
|
||||
Sequence parallel is to support long-sequence modelling such as document-level text understanding and medical imaging.
|
||||
This method is proposed in [Sequence Parallelism: Making 4D Parallelism Possible](https://arxiv.org/abs/2105.13120).
|
||||
You can use specify the mode to be `sequence` to initialize its process group.
|
||||
|
||||
|
||||
```python
|
||||
parallel = dict(
|
||||
tensor=dict(size=4, mode='sequence')
|
||||
)
|
||||
```
|
82
docs/source/en/basics/define_your_config.md
Normal file
82
docs/source/en/basics/define_your_config.md
Normal file
@@ -0,0 +1,82 @@
|
||||
# Define Your Configuration
|
||||
|
||||
Author: Guangyang Lu, Shenggui Li, Siqi Mai
|
||||
|
||||
**Prerequisite:**
|
||||
- [Distributed Training](../concepts/distributed_training.md)
|
||||
- [Colossal-AI Overview](../concepts/colossalai_overview.md)
|
||||
|
||||
|
||||
## Introduction
|
||||
|
||||
In Colossal-AI, a configuration file is required to specify the features the system will inject into the training process.
|
||||
In this tutorial, we will introduce you how to construct your configuration file and how this config file will be used.
|
||||
Using configuration file has several advantages:
|
||||
|
||||
1. You can store your feature configuration and training hyper-parameters in different configuration files
|
||||
2. New features released in the future can be specified in the configuration without code change in the training script
|
||||
|
||||
In this tutorial, we will cover how to define your configuration file.
|
||||
|
||||
## Configuration Definition
|
||||
|
||||
In a configuration file, there are two types of variables. One serves as feature specification and the other serves
|
||||
as hyper-parameters. All feature-related variables are reserved keywords. For example, if you want to use mixed precision
|
||||
training, you need to use the variable name `fp16` in the config file and follow a pre-defined format.
|
||||
|
||||
### Feature Specification
|
||||
|
||||
There is an array of features Colossal-AI provides to speed up training. Each feature is defined by a corresponding field
|
||||
in the config file. In this tutorial, we are not giving the config details for all the features, but rather we are providing
|
||||
an illustration of how to specify a feature. **The details of each feature can be found in its respective tutorial.**
|
||||
|
||||
To illustrate the use of config file, we use mixed precision training as an example here. In order to do so, you need to
|
||||
follow the steps below.
|
||||
|
||||
1. create a configuration file (e.g. `config.py`, the file name can be anything)
|
||||
2. define the mixed precision configuration in the config file. For example, in order to use mixed precision training
|
||||
natively provided by PyTorch, you can just write these lines of code below into your config file.
|
||||
|
||||
```python
|
||||
from colossalai.amp import AMP_TYPE
|
||||
|
||||
fp16 = dict(
|
||||
mode=AMP_TYPE.TORCH
|
||||
)
|
||||
```
|
||||
|
||||
3. Tell Colossal-AI where your config file is when launch the distributed environment. For example, the config file is in
|
||||
the current directory.
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
|
||||
colossalai.launch(config='./config.py', ...)
|
||||
```
|
||||
|
||||
In this way, Colossal-AI knows what features you want to use and will inject this feature during `colossalai.initialize`.
|
||||
|
||||
### Global Hyper-parameters
|
||||
|
||||
Besides feature specification, the config file can also serve as a place to define your training hyper-parameters. This
|
||||
comes handy when you want to perform multiple experiments, each experiment details can be put into a single config file
|
||||
to avoid confusion. These parameters will be stored in the global parallel context and can be accessed in the training script.
|
||||
|
||||
For example, you can specify the batch size in your config file.
|
||||
|
||||
```python
|
||||
BATCH_SIZE = 32
|
||||
```
|
||||
|
||||
After launch, you are able to access your hyper-parameters through global parallel context.
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
from colossalai.core import global_context as gpc
|
||||
|
||||
colossalai.launch(config='./config.py', ...)
|
||||
|
||||
# access your parameter
|
||||
print(gpc.config.BATCH_SIZE)
|
||||
|
||||
```
|
387
docs/source/en/basics/engine_trainer.md
Normal file
387
docs/source/en/basics/engine_trainer.md
Normal file
@@ -0,0 +1,387 @@
|
||||
# Use Engine and Trainer in Training
|
||||
|
||||
Author: Shenggui Li, Siqi Mai
|
||||
|
||||
**Prerequisite:**
|
||||
- [Initialize Features](./initialize_features.md)
|
||||
|
||||
## Introduction
|
||||
|
||||
In this tutorial, you will learn how to use the engine and trainer provided in Colossal-AI to train your model.
|
||||
Before we delve into the details, we would like to first explain the concept of engine and trainer.
|
||||
|
||||
### Engine
|
||||
|
||||
Engine is essentially a wrapper class for model, optimizer and loss function.
|
||||
When we call `colossalai.initialize`, an engine object will be returned, and it has already been equipped with
|
||||
functionalities such as gradient clipping, gradient accumulation and zero optimizer as specified in your configuration file.
|
||||
An engine object will use similar APIs to those of PyTorch training components such that the user has minimum change
|
||||
to their code.
|
||||
|
||||
Below is a table which shows the commonly used APIs for the engine object.
|
||||
|
||||
| Component | Function | PyTorch | Colossal-AI |
|
||||
| ------------------------------------- | --------------------------------------------- | ------------------------------- | -------------------------------------- |
|
||||
| optimizer | Set all gradients to zero before an iteration | optimizer.zero_grad() | engine.zero_grad() |
|
||||
| optimizer | Update the parameters | optimizer.step() | engine.step() |
|
||||
| model | Run a forward pass | outputs = model(inputs) | outputs = engine(inputs) |
|
||||
| criterion | Calculate the loss value | loss = criterion(output, label) | loss = engine.criterion(output, label) |
|
||||
| criterion | Execute back-propagation on the model | loss.backward() | engine.backward(loss) |
|
||||
|
||||
The reason why we need such an engine class is that we can add more functionalities while hiding the implementations in
|
||||
the `colossalai.initialize` function.
|
||||
Imaging we are gonna add a new feature, we can manipulate the model, optimizer, dataloader and loss function in the
|
||||
`colossalai.initialize` function and only expose an engine object to the user.
|
||||
The user only needs to modify their code to the minimum extent by adapting the normal PyTorch APIs to the Colossal-AI
|
||||
engine APIs. In this way, they can enjoy more features for efficient training.
|
||||
|
||||
A normal training iteration using engine can be:
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
|
||||
# build your model, optimizer, criterion, dataloaders
|
||||
...
|
||||
|
||||
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model,
|
||||
optimizer,
|
||||
criterion,
|
||||
train_dataloader,
|
||||
test_dataloader)
|
||||
for img, label in train_dataloader:
|
||||
engine.zero_grad()
|
||||
output = engine(img)
|
||||
loss = engine.criterion(output, label)
|
||||
engine.backward(loss)
|
||||
engine.step()
|
||||
```
|
||||
|
||||
### Trainer
|
||||
|
||||
Trainer is a more high-level wrapper for the user to execute training with fewer lines of code. However, in pursuit of more abstraction, it loses some flexibility compared to engine. The trainer is designed to execute a forward and backward step to perform model weight update. It is easy to create a trainer object by passing the engine object. The trainer has a default value `None` for the argument `schedule`. In most cases, we leave this value to `None` unless we want to use pipeline parallelism. If you wish to explore more about this parameter, you can go to the tutorial on pipeline parallelism.
|
||||
|
||||
```python
|
||||
from colossalai.logging import get_dist_logger
|
||||
from colossalai.trainer import Trainer, hooks
|
||||
|
||||
# build components and initialize with colossalai.initialize
|
||||
...
|
||||
|
||||
# create a logger so that trainer can log on the console
|
||||
logger = get_dist_logger()
|
||||
|
||||
# create a trainer object
|
||||
trainer = Trainer(
|
||||
engine=engine,
|
||||
logger=logger
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
|
||||
In trainer, the user can customize some hooks and attach these hooks to the trainer object. A hook object will execute life-cycle methods periodically based on the training scheme. For example, The `LRSchedulerHook` will execute `lr_scheduler.step()` to update the learning rate of the model during either `after_train_iter` or `after_train_epoch` stages depending on whether the user wants to update the learning rate after each training iteration or only after the entire training epoch. You can store the hook objects in a list and pass it to `trainer.fit` method. `trainer.fit` method will execute training and testing based on your parameters. If `display_process` is True, a progress bar will be displayed on your console to show the training process.
|
||||
|
||||
```python
|
||||
# define the hooks to attach to the trainer
|
||||
hook_list = [
|
||||
hooks.LossHook(),
|
||||
hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=True),
|
||||
hooks.AccuracyHook(accuracy_func=Accuracy()),
|
||||
hooks.LogMetricByEpochHook(logger),
|
||||
]
|
||||
|
||||
# start training
|
||||
trainer.fit(
|
||||
train_dataloader=train_dataloader,
|
||||
epochs=NUM_EPOCHS,
|
||||
test_dataloader=test_dataloader,
|
||||
test_interval=1,
|
||||
hooks=hook_list,
|
||||
display_progress=True
|
||||
)
|
||||
```
|
||||
|
||||
If you want to customize your own hook class, you can inherit `hooks.BaseHook` and override the life-cycle methods of your interest. A dummy example to demonstrate how to create a simple log message hook is provided below for your reference.
|
||||
|
||||
```python
|
||||
from colossalai.logging import get_dist_logger
|
||||
from colossalai.trainer import hooks
|
||||
|
||||
class LogMessageHook(hooks.BaseHook):
|
||||
|
||||
def __init__(self, priority=10):
|
||||
self._logger = get_dist_logger()
|
||||
|
||||
def before_train(self, trainer):
|
||||
self._logger.info('training starts')
|
||||
|
||||
def after_train(self, trainer):
|
||||
self._logger.info('training finished')
|
||||
|
||||
|
||||
...
|
||||
|
||||
# then in your training script
|
||||
hook_list.append(LogMessageHook())
|
||||
```
|
||||
|
||||
|
||||
|
||||
In the sections below, I will guide you through the steps required to train a ResNet model with both engine and trainer.
|
||||
|
||||
|
||||
|
||||
## Explain with ResNet
|
||||
|
||||
### Overview
|
||||
|
||||
In this section we will cover:
|
||||
|
||||
1. Use an engine object to train a ResNet34 model on CIFAR10 dataset
|
||||
2. Use a trainer object to train a ResNet34 model on CIFAR10 dataset
|
||||
|
||||
The project structure will be like:
|
||||
|
||||
```bash
|
||||
-- config.py
|
||||
-- run_resnet_cifar10_with_engine.py
|
||||
-- run_resnet_cifar10_with_trainer.py
|
||||
```
|
||||
|
||||
Steps 1-4 below are commonly used regardless of using engine or trainer. Thus, steps 1-4 + step 5 will be your `run_resnet_cifar10_with_engine.py` and steps 1-4 + step 6 will form `run_resnet_cifar10_with_trainer.py`.
|
||||
|
||||
### Hands-on Practice
|
||||
|
||||
#### Step 1. Create a Config File
|
||||
|
||||
In your project folder, create a `config.py`. This file is to specify some features you may want to use to train your model. A sample config file is as below:
|
||||
|
||||
```python
|
||||
from colossalai.amp import AMP_TYPE
|
||||
|
||||
BATCH_SIZE = 128
|
||||
NUM_EPOCHS = 200
|
||||
|
||||
fp16=dict(
|
||||
mode=AMP_TYPE.TORCH
|
||||
)
|
||||
```
|
||||
|
||||
In this config file, we specify that we want to use batch size 128 per GPU and run for 200 epochs. These two parameters are exposed by `gpc.config`. For example, you can use `gpc.config.BATCH_SIZE` to access the value you store in your config file. The `fp16` configuration tells `colossalai.initialize` to use mixed precision training provided by PyTorch to train the model with better speed and lower memory consumption.
|
||||
|
||||
#### Step 2. Initialize Distributed Environment
|
||||
|
||||
We need to initialize the distributed training environment. This has been introduced in the tutorial on how to
|
||||
[launch Colossal-AI](./launch_colossalai.md). For this demostration, we use `launch_from_torch` and PyTorch launch utility.
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
|
||||
# ./config.py refers to the config file we just created in step 1
|
||||
colossalai.launch_from_torch(config='./config.py')
|
||||
```
|
||||
|
||||
#### Step 3. Create all the training components
|
||||
|
||||
In this step, we can create all the components used for training. These components include:
|
||||
|
||||
1. Model
|
||||
2. Optimizer
|
||||
3. Criterion/loss function
|
||||
4. Training/Testing dataloaders
|
||||
5. Learning rate Scheduler
|
||||
6. Logger
|
||||
|
||||
|
||||
|
||||
To build these components, you need to import the following modules:
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
from colossalai.logging import get_dist_logger
|
||||
import torch
|
||||
import os
|
||||
from colossalai.core import global_context as gpc
|
||||
from colossalai.utils import get_dataloader
|
||||
from torchvision import transforms
|
||||
from colossalai.nn.lr_scheduler import CosineAnnealingLR
|
||||
from torchvision.datasets import CIFAR10
|
||||
from torchvision.models import resnet34
|
||||
```
|
||||
|
||||
|
||||
|
||||
Then build your components in the same way as how to normally build them in your PyTorch scripts. In the script below, we set the root path for CIFAR10 dataset as an environment variable `DATA`. You can change it to any path you like, for example, you can change `root=Path(os.environ['DATA'])` to `root='./data'` so that there is no need to set the environment variable.
|
||||
|
||||
```python
|
||||
# build logger
|
||||
logger = get_dist_logger()
|
||||
|
||||
# build resnet
|
||||
model = resnet34(num_classes=10)
|
||||
|
||||
# build datasets
|
||||
train_dataset = CIFAR10(
|
||||
root='./data',
|
||||
download=True,
|
||||
transform=transforms.Compose(
|
||||
[
|
||||
transforms.RandomCrop(size=32, padding=4),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[
|
||||
0.2023, 0.1994, 0.2010]),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
test_dataset = CIFAR10(
|
||||
root='./data',
|
||||
train=False,
|
||||
transform=transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[
|
||||
0.2023, 0.1994, 0.2010]),
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
# build dataloaders
|
||||
train_dataloader = get_dataloader(dataset=train_dataset,
|
||||
shuffle=True,
|
||||
batch_size=gpc.config.BATCH_SIZE,
|
||||
num_workers=1,
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
test_dataloader = get_dataloader(dataset=test_dataset,
|
||||
add_sampler=False,
|
||||
batch_size=gpc.config.BATCH_SIZE,
|
||||
num_workers=1,
|
||||
pin_memory=True,
|
||||
)
|
||||
|
||||
# build criterion
|
||||
criterion = torch.nn.CrossEntropyLoss()
|
||||
|
||||
# optimizer
|
||||
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
|
||||
|
||||
# lr_scheduler
|
||||
lr_scheduler = CosineAnnealingLR(optimizer, total_steps=gpc.config.NUM_EPOCHS)
|
||||
```
|
||||
|
||||
#### Step 4. Initialize with Colossal-AI
|
||||
|
||||
Next, the essential step is to obtain the engine class by calling `colossalai.initialize`. As stated in `config.py`, we will be using mixed precision training for training ResNet34 model. `colossalai.initialize` will automatically check your config file and assign relevant features to your training components. In this way, our engine object has already been able to train with mixed precision, but you do not have to explicitly take care of it.
|
||||
|
||||
```python
|
||||
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model,
|
||||
optimizer,
|
||||
criterion,
|
||||
train_dataloader,
|
||||
test_dataloader,
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
|
||||
#### Step 5. Train with engine
|
||||
|
||||
With all the training components ready, we can train ResNet34 just like how to normally deal with PyTorch training.
|
||||
|
||||
```python
|
||||
for epoch in range(gpc.config.NUM_EPOCHS):
|
||||
# execute a training iteration
|
||||
engine.train()
|
||||
for img, label in train_dataloader:
|
||||
img = img.cuda()
|
||||
label = label.cuda()
|
||||
|
||||
# set gradients to zero
|
||||
engine.zero_grad()
|
||||
|
||||
# run forward pass
|
||||
output = engine(img)
|
||||
|
||||
# compute loss value and run backward pass
|
||||
train_loss = engine.criterion(output, label)
|
||||
engine.backward(train_loss)
|
||||
|
||||
# update parameters
|
||||
engine.step()
|
||||
|
||||
# update learning rate
|
||||
lr_scheduler.step()
|
||||
|
||||
# execute a testing iteration
|
||||
engine.eval()
|
||||
correct = 0
|
||||
total = 0
|
||||
for img, label in test_dataloader:
|
||||
img = img.cuda()
|
||||
label = label.cuda()
|
||||
|
||||
# run prediction without back-propagation
|
||||
with torch.no_grad():
|
||||
output = engine(img)
|
||||
test_loss = engine.criterion(output, label)
|
||||
|
||||
# compute the number of correct prediction
|
||||
pred = torch.argmax(output, dim=-1)
|
||||
correct += torch.sum(pred == label)
|
||||
total += img.size(0)
|
||||
|
||||
logger.info(
|
||||
f"Epoch {epoch} - train loss: {train_loss:.5}, test loss: {test_loss:.5}, acc: {correct / total:.5}, lr: {lr_scheduler.get_last_lr()[0]:.5g}", ranks=[0])
|
||||
```
|
||||
|
||||
#### Step 6. Train with trainer
|
||||
|
||||
If you wish to train with a trainer object, you can follow the code snippet below:
|
||||
|
||||
```python
|
||||
from colossalai.nn.metric import Accuracy
|
||||
from colossalai.trainer import Trainer, hooks
|
||||
|
||||
|
||||
# create a trainer object
|
||||
trainer = Trainer(
|
||||
engine=engine,
|
||||
logger=logger
|
||||
)
|
||||
|
||||
# define the hooks to attach to the trainer
|
||||
hook_list = [
|
||||
hooks.LossHook(),
|
||||
hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=True),
|
||||
hooks.AccuracyHook(accuracy_func=Accuracy()),
|
||||
hooks.LogMetricByEpochHook(logger),
|
||||
hooks.LogMemoryByEpochHook(logger)
|
||||
]
|
||||
|
||||
# start training
|
||||
# run testing every 1 epoch
|
||||
trainer.fit(
|
||||
train_dataloader=train_dataloader,
|
||||
epochs=gpc.config.NUM_EPOCHS,
|
||||
test_dataloader=test_dataloader,
|
||||
test_interval=1,
|
||||
hooks=hook_list,
|
||||
display_progress=True
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
|
||||
#### Step 7. Start Distributed Training
|
||||
|
||||
Lastly, we can invoke the scripts using the distributed launcher provided by PyTorch as we used `launch_from_torch` in Step 2. You need to replace `<num_gpus>` with the number of GPUs available on your machine. This number can be 1 if you only want to use 1 GPU. If you wish to use other launchers, you can refer to the tutorial on How to Launch Colossal-AI.
|
||||
|
||||
```bash
|
||||
# with engine
|
||||
python -m torch.distributed.launch --nproc_per_node <num_gpus> --master_addr localhost --master_port 29500 run_resnet_cifar10_with_engine.py
|
||||
# with trainer
|
||||
python -m torch.distributed.launch --nproc_per_node <num_gpus> --master_addr localhost --master_port 29500 run_resnet_cifar10_with_trainer.py
|
||||
```
|
49
docs/source/en/basics/initialize_features.md
Normal file
49
docs/source/en/basics/initialize_features.md
Normal file
@@ -0,0 +1,49 @@
|
||||
# Initialize Features
|
||||
|
||||
Author: Shenggui Li, Siqi Mai
|
||||
|
||||
**Prerequisite:**
|
||||
- [Distributed Training](../concepts/distributed_training.md)
|
||||
- [Colossal-AI Overview](../concepts/colossalai_overview.md)
|
||||
|
||||
## Introduction
|
||||
|
||||
In this tutorial, we will cover the use of `colossalai.initialize` which injects features into your training components
|
||||
(e.g. model, optimizer, dataloader) seamlessly. Calling `colossalai.initialize` is the standard procedure before you run
|
||||
into your training loops.
|
||||
|
||||
In the section below, I will cover how `colossalai.initialize` works and what we should take note of.
|
||||
|
||||
## Usage
|
||||
|
||||
In a typical workflow, we will launch distributed environment at the beginning of our training script.
|
||||
Afterwards, we will instantiate our objects such as model, optimizer, loss function, dataloader etc. At this moment, `colossalai.initialize`
|
||||
can come in to inject features into these objects. A pseudo-code example is like below:
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
import torch
|
||||
...
|
||||
|
||||
|
||||
# launch distributed environment
|
||||
colossalai.launch(config='./config.py', ...)
|
||||
|
||||
# create your objects
|
||||
model = MyModel()
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
|
||||
criterion = torch.nn.CrossEntropyLoss()
|
||||
train_dataloader = MyTrainDataloader()
|
||||
test_dataloader = MyTrainDataloader()
|
||||
|
||||
# initialize features
|
||||
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(model,
|
||||
optimizer,
|
||||
criterion,
|
||||
train_dataloader,
|
||||
test_dataloader)
|
||||
```
|
||||
|
||||
The `colossalai.initialize` function will return an `Engine` object. The engine object is a wrapper
|
||||
for model, optimizer and loss function. **The engine object will run with features specified in the config file.**
|
||||
More details about the engine can be found in the [Use Engine and Trainer in Training](./engine_trainer.md).
|
232
docs/source/en/basics/launch_colossalai.md
Normal file
232
docs/source/en/basics/launch_colossalai.md
Normal file
@@ -0,0 +1,232 @@
|
||||
# Launch Colossal-AI
|
||||
|
||||
Author: Chuanrui Wang, Shenggui Li, Siqi Mai
|
||||
|
||||
**Prerequisite:**
|
||||
- [Distributed Training](../concepts/distributed_training.md)
|
||||
- [Colossal-AI Overview](../concepts/colossalai_overview.md)
|
||||
|
||||
|
||||
## Introduction
|
||||
|
||||
As mentioned in the previous tutorials stated in the prerequisite, you need to initialize the distributed environment
|
||||
for Colossal-AI after your config file is prepared.
|
||||
We call this process `launch`.
|
||||
In this tutorial, you will learn how to launch Colossal-AI on your server, be it a small one or big one.
|
||||
|
||||
In Colossal-AI, we provided several launch methods to initialize the distributed backend.
|
||||
In most cases, you can use `colossalai.launch` and `colossalai.get_default_parser` to pass the
|
||||
parameters via command line.
|
||||
If you happen to use launchers such as SLURM, OpenMPI and PyTorch launch utility,
|
||||
we also provide several launching helper methods to access the rank and world size from the environment variables
|
||||
set by these launchers directly for your convenience.
|
||||
|
||||
In this tutorial we will cover how to launch Colossal-AI to initialize the distributed backends:
|
||||
- Launch with `colossalai.launch`
|
||||
- Launch with Colossal-AI CLI
|
||||
- Launch with SLURM
|
||||
- Launch with OpenMPI
|
||||
|
||||
## Launch Distributed Environment
|
||||
|
||||
In order to launch Colossal-AI, we need two types of arguments:
|
||||
1. config file
|
||||
2. distributed settings
|
||||
|
||||
The config file is always required regardless of the launch method but distributed settings can vary. The config file
|
||||
can be a path to the configuration file or a Python dictionary. The distributed settings can be passed via command line
|
||||
or multi-process launchers.
|
||||
|
||||
### Command Line Parser
|
||||
|
||||
Before we jump to `launch`, we firstly need to understand what parameters we need for initialization.
|
||||
As stated in the `Basic Concepts in Distributed Training` section of [Distributed Training](../concepts/distributed_training.md),
|
||||
the important parameters are:
|
||||
|
||||
1. host
|
||||
2. port
|
||||
3. rank
|
||||
4. world_size
|
||||
5. backend
|
||||
|
||||
In Colossal-AI, we provided a command line parser which has added these arguments in advance. You can get this parser by calling
|
||||
`colossalai.get_default_parser()`. This parser is usually used with `colossalai.launch`.
|
||||
|
||||
```python
|
||||
# add these lines in your train.py
|
||||
import colossalai
|
||||
|
||||
# get default parser
|
||||
parser = colossalai.get_default_parser()
|
||||
|
||||
# if you want to add your own arguments
|
||||
parser.add_argument(...)
|
||||
|
||||
# parse arguments
|
||||
args = parser.parse_args()
|
||||
```
|
||||
|
||||
Then in your terminal, you can pass in these arguments:
|
||||
```shell
|
||||
|
||||
python train.py --host <host> --rank <rank> --world_size <world_size> --port <port> --backend <backend>
|
||||
```
|
||||
|
||||
`backend` is optional and the default value is `nccl`.
|
||||
|
||||
### Native Launch
|
||||
|
||||
To initialize the distributed environment, we provided a general `colossalai.launch` API. The `colossalai.launch` function takes in the parameters
|
||||
listed above and create a default process group in the communication network. This function is often used with the default
|
||||
parser for convenience.
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
|
||||
# parse arguments
|
||||
args = colossalai.get_default_parser().parse_args()
|
||||
|
||||
# launch distributed environment
|
||||
colossalai.launch(config=<CONFIG>,
|
||||
rank=args.rank,
|
||||
world_size=args.world_size,
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
backend=args.backend
|
||||
)
|
||||
|
||||
```
|
||||
|
||||
|
||||
### Launch with Colossal-AI CLI
|
||||
|
||||
To enable easy launching on both single or multi nodes, we have implemented a launcher for Colossal-AI. This launcher is
|
||||
a wrapper of the torch distributed launch utility but enhanced with the capability of launching multi-node jobs easily.
|
||||
|
||||
First, we need to set the launch method in our code. As this is a wrapper of the torch distributed launch utility, we will
|
||||
use `colossalai.launch_from_torch`. The arguments required for distributed environment such as rank, world size, host and port are all set by the PyTorch
|
||||
launcher and can be read from the environment variable directly.
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
|
||||
colossalai.launch_from_torch(
|
||||
config=<CONFIG>,
|
||||
)
|
||||
```
|
||||
|
||||
Next, we can easily start multiple processes with `colossalai run` in your terminal. Below is an example to run the code
|
||||
on a single node with 4 GPUs. You can change the number of GPUs by `nproc_per_node` and the default port by `master_port`.
|
||||
|
||||
```shell
|
||||
# run on the local node with 4 GPUs (default port: 29500)
|
||||
colossalai run --nproc_per_node 4 train.py
|
||||
|
||||
# run on the local node with 4 GPUs with a different port
|
||||
colossalai run --nproc_per_node 4 --master_port 29505 test.py
|
||||
```
|
||||
|
||||
If you are in a cluster and want to launch multi-node training, the CLI can help you start processes on different nodes
|
||||
with one simple command. There are two ways you can launch multi-node jobs.
|
||||
|
||||
- Run with `--hosts`
|
||||
|
||||
This is suitable when you only have a few nodes. Let's say I have two nodes, namely `host1` and `host2`, I can start
|
||||
multi-node training with the following command. Compared to single-node training, you must specify the `master_addr`
|
||||
option, which is auto-set to localhost if running on a single node only.
|
||||
|
||||
:::caution
|
||||
|
||||
`master_addr` cannot be localhost when running on multiple nodes, it should be the hostname or IP address of a node.
|
||||
|
||||
:::
|
||||
|
||||
```shell
|
||||
# run on these two nodes
|
||||
colossalai run --nproc_per_node 4 --host host1,host2 --master_addr host1 test.py
|
||||
```
|
||||
- Run with `--hostfile`
|
||||
|
||||
This method is suitable when you have a lot of nodes. The host file is a simple text file listing the available nodes.
|
||||
The list of nodes is commonly provided by cluster managers such as SLURM and PBS Pro. For example, you can get the list
|
||||
of nodes allocated to you via the environment variable `SLURM_NODELIST` in SLURM and `PBS_NODEFILE` in PBS Pro.
|
||||
Just do `echo $SLURM_NODELIST` or `cat $PBS_NODEFILE` to check it out. If you do not have such cluster managers, you can
|
||||
manually create one for your own use.
|
||||
|
||||
The host file given to Colossal-AI launcher must be in the following format where each line is the host name of a node.
|
||||
|
||||
```text
|
||||
host1
|
||||
host2
|
||||
```
|
||||
|
||||
With the host file ready, we can launch multi-node jobs with the following commands. Just like using `--host`, you also
|
||||
need to specify the `master_addr` option. Some extra options are provided for `--hostfile` as listed below:
|
||||
|
||||
- `--include`: specify the hosts to include for multi-node jobs. For example, if your host file has 8 nodes, but you
|
||||
happen to only want to run on 6 nodes instead, you can add `--include host1,host2,host3,...,host6` so that the job will only
|
||||
be launcher on the 6 nodes.
|
||||
- `--exclude`: specify the hosts to exclude for multi-node jobs. This is useful when some nodes are faulty. For example,
|
||||
if host1 GPU has some problems and you do not wish to run on host1 but all other nodes, you can add `--exclude host1` so that
|
||||
the job will only be launched on the remaining nodes.
|
||||
|
||||
```shell
|
||||
# run with a hostfile
|
||||
colossalai run --nproc_per_node 4 --hostfile ./hostfile --master_addr host1 test.py
|
||||
|
||||
# only include certain hosts to execute commands
|
||||
# this is used to manually select nodes to run
|
||||
colossalai run --nproc_per_node 4 --hostfile ./hostfile --master_addr host1 --include host1 test.py
|
||||
|
||||
# exclude certain hosts to execute commands
|
||||
# this can be used when certain nodes are faulty
|
||||
colossalai run --nproc_per_node 4 --hostfile ./hostfile --master_addr host1 --exclude host2 test.py
|
||||
```
|
||||
|
||||
### Launch with SLURM
|
||||
|
||||
If you are on a system managed by the SLURM scheduler, you can also rely on the `srun` launcher to kickstart your Colossal-AI scripts.
|
||||
We provided the helper function `launch_from_slurm` for compatibility with the SLURM scheduler.
|
||||
`launch_from_slurm` will automatically read the rank and world size from the environment variables `SLURM_PROCID` and `SLURM_NPROCS` respectively
|
||||
and use them to start the distributed backend.
|
||||
Do this in your training script:
|
||||
|
||||
```python
|
||||
import colossalai
|
||||
|
||||
colossalai.launch_from_slurm(
|
||||
config=<CONFIG>,
|
||||
host=args.host,
|
||||
port=args.port
|
||||
)
|
||||
```
|
||||
|
||||
You can initialize the distributed environment by using this command in terminal.
|
||||
|
||||
```bash
|
||||
srun python train.py --host <master_node> --port 29500
|
||||
```
|
||||
|
||||
### Launch with OpenMPI
|
||||
If you are more familiar with OpenMPI, you can use `launch_from_openmpi` instead.
|
||||
`launch_from_openmpi` will automatically read the local rank, global rank and world size from the environment variables
|
||||
`OMPI_COMM_WORLD_LOCAL_RANK`, `MPI_COMM_WORLD_RANK` and `OMPI_COMM_WORLD_SIZE` respectively and
|
||||
use them to start the distributed backend.
|
||||
|
||||
Do this in your train.py:
|
||||
```python
|
||||
colossalai.launch_from_openmpi(
|
||||
config=<CONFIG>,
|
||||
host=args.host,
|
||||
port=args.port
|
||||
)
|
||||
```
|
||||
|
||||
A sample command to launch multiple processes with OpenMPI would be:
|
||||
|
||||
```bash
|
||||
mpirun --hostfile <my_hostfile> -np <num_process> python train.py --host <node name or ip> --port 29500
|
||||
```
|
||||
|
||||
- --hostfile: use this option to specify a list of hosts on which to run
|
||||
- --np: set the number of processes (GPUs) to launch in total. For example, if --np 4, 4 python processes will be initialized to run train.py.
|
61
docs/source/en/basics/model_checkpoint.md
Normal file
61
docs/source/en/basics/model_checkpoint.md
Normal file
@@ -0,0 +1,61 @@
|
||||
# Model Checkpoint
|
||||
|
||||
Author : Guangyang Lu
|
||||
|
||||
**Prerequisite:**
|
||||
- [Launch Colossal-AI](./launch_colossalai.md)
|
||||
- [Initialize Colossal-AI](./initialize_features.md)
|
||||
|
||||
**Example Code:**
|
||||
- [ColossalAI-Examples Model Checkpoint](https://github.com/hpcaitech/ColossalAI-Examples/tree/main/utils/checkpoint)
|
||||
|
||||
**This function is experiential.**
|
||||
|
||||
## Introduction
|
||||
|
||||
In this tutorial, you will learn how to save and load model checkpoints.
|
||||
|
||||
To leverage the power of parallel strategies in Colossal-AI, modifications to models and tensors are needed, for which you cannot directly use `torch.save` or `torch.load` to save or load model checkpoints. Therefore, we have provided you with the API to achieve the same thing.
|
||||
|
||||
Moreover, when loading, you are not demanded to use the same parallel strategy as saving.
|
||||
|
||||
## How to use
|
||||
|
||||
### Save
|
||||
|
||||
There are two ways to train a model in Colossal-AI, by engine or by trainer.
|
||||
**Be aware that we only save the `state_dict`.** Therefore, when loading the checkpoints, you need to define the model first.
|
||||
|
||||
#### Save when using engine
|
||||
|
||||
```python
|
||||
from colossalai.utils import save_checkpoint
|
||||
model = ...
|
||||
engine, _, _, _ = colossalai.initialize(model=model, ...)
|
||||
for epoch in range(num_epochs):
|
||||
... # do some training
|
||||
save_checkpoint('xxx.pt', epoch, model)
|
||||
```
|
||||
|
||||
#### Save when using trainer
|
||||
```python
|
||||
from colossalai.trainer import Trainer, hooks
|
||||
model = ...
|
||||
engine, _, _, _ = colossalai.initialize(model=model, ...)
|
||||
trainer = Trainer(engine, ...)
|
||||
hook_list = [
|
||||
hooks.SaveCheckpointHook(1, 'xxx.pt', model)
|
||||
...]
|
||||
|
||||
trainer.fit(...
|
||||
hook=hook_list)
|
||||
```
|
||||
|
||||
### Load
|
||||
|
||||
```python
|
||||
from colossalai.utils import load_checkpoint
|
||||
model = ...
|
||||
load_checkpoint('xxx.pt', model)
|
||||
... # train or test
|
||||
```
|
Reference in New Issue
Block a user