2021-10-28 16:21:23 +00:00
|
|
|
#!/usr/bin/env python
|
|
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
|
|
|
|
import pytest
|
|
|
|
import torch
|
|
|
|
import torch.distributed as dist
|
2023-04-06 06:51:35 +00:00
|
|
|
|
2023-09-11 08:24:28 +00:00
|
|
|
from colossalai.context.parallel_mode import ParallelMode
|
|
|
|
from colossalai.core import global_context as gpc
|
|
|
|
from colossalai.initialize import launch
|
|
|
|
from colossalai.legacy.communication import (
|
2023-04-06 06:51:35 +00:00
|
|
|
recv_backward,
|
|
|
|
recv_forward,
|
|
|
|
recv_obj_meta,
|
|
|
|
send_backward,
|
|
|
|
send_backward_recv_forward,
|
|
|
|
send_forward,
|
|
|
|
send_forward_recv_backward,
|
|
|
|
send_obj_meta,
|
|
|
|
)
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
from colossalai.logging import get_dist_logger
|
2023-04-06 06:51:35 +00:00
|
|
|
from colossalai.testing import rerun_if_address_is_in_use, spawn
|
|
|
|
from colossalai.utils import get_current_device
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2022-03-17 07:44:17 +00:00
|
|
|
BATCH_SIZE = 4
|
|
|
|
SEQ_LENGTH = 2
|
|
|
|
HIDDEN_SIZE = 16
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2022-03-17 07:44:17 +00:00
|
|
|
CONFIG = dict(parallel=dict(pipeline=dict(size=4), tensor=dict(size=1, mode=None)), seed=1024)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
|
|
|
|
def check_equal(A, B):
|
|
|
|
return torch.allclose(A, B, rtol=1e-5, atol=1e-3)
|
|
|
|
|
|
|
|
|
|
|
|
def check_forward(output_tensor, rank, logger):
|
|
|
|
dist.barrier()
|
|
|
|
if gpc.is_first_rank(ParallelMode.PIPELINE):
|
|
|
|
tensor = output_tensor.clone()
|
|
|
|
else:
|
|
|
|
tensor = recv_forward(output_tensor.shape)
|
2022-03-17 07:44:17 +00:00
|
|
|
logger.info('Rank {} received forward. Correct tensor: {}'.format(rank, check_equal(tensor, output_tensor)))
|
2021-10-28 16:21:23 +00:00
|
|
|
if not gpc.is_last_rank(ParallelMode.PIPELINE):
|
|
|
|
send_forward(tensor)
|
|
|
|
logger.info('Rank {} sent forward.'.format(rank))
|
|
|
|
|
|
|
|
|
|
|
|
def check_backward(output_grad, rank, logger):
|
|
|
|
dist.barrier()
|
|
|
|
if gpc.is_last_rank(ParallelMode.PIPELINE):
|
|
|
|
grad = output_grad.clone()
|
|
|
|
else:
|
|
|
|
grad = recv_backward(output_grad.shape)
|
2022-03-17 07:44:17 +00:00
|
|
|
logger.info('Rank {} received backward. Correct grad: {}'.format(rank, check_equal(grad, output_grad)))
|
2021-10-28 16:21:23 +00:00
|
|
|
if not gpc.is_first_rank(ParallelMode.PIPELINE):
|
|
|
|
send_backward(grad)
|
|
|
|
logger.info('Rank {} sent backward.'.format(rank))
|
|
|
|
|
|
|
|
|
|
|
|
def check_forward_backward(output_tensor, output_grad, rank, logger):
|
|
|
|
dist.barrier()
|
|
|
|
if not gpc.is_first_rank(ParallelMode.PIPELINE):
|
|
|
|
tensor = send_backward_recv_forward(output_grad, output_tensor.shape)
|
2022-03-17 07:44:17 +00:00
|
|
|
logger.info('Rank {} sent backward received forward. Correct tensor: {}'.format(
|
|
|
|
rank, check_equal(tensor, output_tensor)))
|
2021-10-28 16:21:23 +00:00
|
|
|
if not gpc.is_last_rank(ParallelMode.PIPELINE):
|
|
|
|
grad = send_forward_recv_backward(output_tensor, output_grad.shape)
|
2022-03-17 07:44:17 +00:00
|
|
|
logger.info('Rank {} sent forward received backward. Correct grad: {}'.format(
|
|
|
|
rank, check_equal(grad, output_grad)))
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
|
2022-03-17 07:44:17 +00:00
|
|
|
def check_comm(size, rank, prev_rank, next_rank, logger):
|
2021-10-28 16:21:23 +00:00
|
|
|
dtype = torch.float32
|
|
|
|
device = get_current_device()
|
|
|
|
tensor_shape = (BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE)
|
|
|
|
grad_shape = (BATCH_SIZE, SEQ_LENGTH, HIDDEN_SIZE)
|
|
|
|
tensor = torch.randn(tensor_shape, dtype=dtype, device=device)
|
|
|
|
dist.all_reduce(tensor)
|
|
|
|
grad = torch.randn(grad_shape, dtype=dtype, device=device)
|
|
|
|
dist.all_reduce(grad)
|
|
|
|
check_forward(tensor, rank, logger)
|
|
|
|
check_backward(grad, rank, logger)
|
|
|
|
check_forward_backward(tensor, grad, rank, logger)
|
|
|
|
|
|
|
|
|
2021-12-29 15:32:10 +00:00
|
|
|
def run_check(rank, world_size, port):
|
2022-03-17 07:44:17 +00:00
|
|
|
launch(config=CONFIG, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
|
Develop/experiments (#59)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
* Split conv2d, class token, positional embedding in 2d, Fix random number in ddp
Fix convergence in cifar10, Imagenet1000
* Integrate 1d tensor parallel in Colossal-AI (#39)
* fixed 1D and 2D convergence (#38)
* optimized 2D operations
* fixed 1D ViT convergence problem
* Feature/ddp (#49)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* support torch ddp
* fix loss accumulation
* add log for ddp
* change seed
* modify timing hook
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* Feature/pipeline (#40)
* remove redundancy func in setup (#19) (#20)
* use env to control the language of doc (#24) (#25)
* Support TP-compatible Torch AMP and Update trainer API (#27)
* Add gradient accumulation, fix lr scheduler
* fix FP16 optimizer and adapted torch amp with tensor parallel (#18)
* fixed bugs in compatibility between torch amp and tensor parallel and performed some minor fixes
* fixed trainer
* Revert "fixed trainer"
This reverts commit 2e0b0b76990e8d4e337add483d878c0f61cf5097.
* improved consistency between trainer, engine and schedule (#23)
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
* add an example of ViT-B/16 and remove w_norm clipping in LAMB (#29)
* add explanation for ViT example (#35) (#36)
* optimize communication of pipeline parallel
* fix grad clip for pipeline
Co-authored-by: Frank Lee <somerlee.9@gmail.com>
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
* optimized 3d layer to fix slow computation ; tested imagenet performance with 3d; reworked lr_scheduler config definition; fixed launch args; fixed some printing issues; simplified apis of 3d layers (#51)
* Update 2.5d layer code to get a similar accuracy on imagenet-1k dataset
* update api for better usability (#58)
update api for better usability
Co-authored-by: 1SAA <c2h214748@gmail.com>
Co-authored-by: ver217 <lhx0217@gmail.com>
Co-authored-by: puck_WCR <46049915+WANG-CR@users.noreply.github.com>
Co-authored-by: binmakeswell <binmakeswell@gmail.com>
Co-authored-by: アマデウス <kurisusnowdeng@users.noreply.github.com>
Co-authored-by: BoxiangW <45734921+BoxiangW@users.noreply.github.com>
2021-12-09 07:08:29 +00:00
|
|
|
logger = get_dist_logger()
|
2021-10-28 16:21:23 +00:00
|
|
|
rank = gpc.get_global_rank()
|
|
|
|
prev_rank = gpc.get_prev_global_rank(ParallelMode.PIPELINE)
|
|
|
|
next_rank = gpc.get_next_global_rank(ParallelMode.PIPELINE)
|
2022-03-17 07:44:17 +00:00
|
|
|
logger.info('Rank {0}: prev rank {1}, next rank {2}'.format(rank, prev_rank, next_rank))
|
2023-05-11 08:30:58 +00:00
|
|
|
logger.info('Distributed environment is initialized.')
|
2021-10-28 16:21:23 +00:00
|
|
|
|
2021-12-30 07:56:46 +00:00
|
|
|
check_comm(world_size, rank, prev_rank, next_rank, logger)
|
2021-12-16 02:32:08 +00:00
|
|
|
gpc.destroy()
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.dist
|
2023-04-06 06:51:35 +00:00
|
|
|
@rerun_if_address_is_in_use()
|
2021-12-16 02:32:08 +00:00
|
|
|
def test_p2p():
|
|
|
|
world_size = 4
|
2023-04-06 06:51:35 +00:00
|
|
|
spawn(run_check, world_size)
|
2021-10-28 16:21:23 +00:00
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
2021-12-16 02:32:08 +00:00
|
|
|
test_p2p()
|