github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-05-03 17:37:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

Support TP-compatible Torch AMP and Update trainer API (#27)

Browse Source 
* 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 2e0b0b7699.

* 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>
This commit is contained in:
Frank Lee
2021-11-18 19:45:06 +08:00
committed by GitHub
parent 2b05de4c64
commit 3defa32aee
80 changed files with 2194 additions and 1584 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_models/test_vision_transformer/configs/vit_2d.py
View File

@@ -102,6 +102,6 @@ parallel = dict(
tensor=dict(size=4, mode='2d'),
)
lr_scheduler = dict(type='LinearWarmupLR', warmup_epochs=5)
num_epochs = 60
lr_scheduler = dict(type='LinearWarmupLR', warmup_steps=5, total_steps=num_epochs)

11
tests/test_models/test_vision_transformer/configs/vit_2p5d.py
View File

@@ -125,13 +125,6 @@ parallel = dict(
tensor=dict(size=4, depth=1, mode='2.5d'),
)
lr_scheduler = dict(
type='LinearWarmupLR',
warmup_epochs=5
)
schedule = dict(
num_microbatches=8
)
num_epochs = 60
lr_scheduler = dict(type='LinearWarmupLR', warmup_steps=5, total_steps=num_epochs)

19
tests/test_models/test_vision_transformer/configs/vit_3d.py
View File

@@ -116,9 +116,14 @@ hooks = [
weight_parallel_mode=ParallelMode.PARALLEL_3D_WEIGHT,
),
dict(type='LossHook'),
# dict(type='TensorboardHook', log_dir='./tfb_logs'),
# dict(type='SaveCheckpointHook', interval=5, checkpoint_dir='./ckpt'),
# dict(type='LoadCheckpointHook', epoch=20, checkpoint_dir='./ckpt')
dict(
type='LRSchedulerHook',
by_epoch=True,
lr_scheduler_cfg=dict(
type='LinearWarmupLR',
warmup_steps=5
)
),
]
parallel = dict(
@@ -127,12 +132,4 @@ parallel = dict(
tensor=dict(mode='3d', size=8),
)
# fp16 = dict(mode=AMP_TYPE.PARALLEL, initial_scale=2 ** 6)
lr_scheduler = dict(type='LinearWarmupLR', warmup_epochs=5)
# schedule = dict(num_microbatches=4)
num_epochs = 60
seed = 42

39
tests/test_models/test_vision_transformer/test_vit_2d/test_vit_2d.py
View File

@@ -7,23 +7,25 @@ import pytest
import torch.autograd
import colossalai
from colossalai.builder import build_lr_scheduler
from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.engine import Engine
from colossalai.logging import get_global_dist_logger
from colossalai.nn.layer._parallel_utilities import _gather
CONFIG_PATH = Path(__file__).parent.parent.joinpath('configs/vit_2d.py')
def eval(engine):
def eval(engine, test_dataloader):
engine.eval()
accumulated_loss = 0
correct_sum = 0
total_sum = 0
num_steps = len(test_dataloader)
data_iter = iter(test_dataloader)
for i in range(engine.schedule.num_steps):
output, label, loss = engine.step()
for i in range(num_steps):
output, label, loss = engine.step(data_iter)
accumulated_loss += loss.detach().cpu().numpy()
output = _gather(
@@ -40,18 +42,21 @@ def eval(engine):
correct = torch.sum(label[0] == output)
correct_sum += correct
total_sum += label[0].size(0)
avg_loss = accumulated_loss / engine.schedule.num_steps
avg_loss = accumulated_loss / num_steps
return correct_sum, total_sum, avg_loss
def train(engine):
def train(engine, train_dataloader, lr_scheduler):
engine.train()
accumulated_loss = 0
num_steps = len(train_dataloader)
data_iter = iter(train_dataloader)
for i in range(engine.schedule.num_steps):
output, label, loss = engine.step()
for i in range(num_steps):
output, label, loss = engine.step(data_iter)
accumulated_loss += loss.detach().cpu().numpy()
avg_loss = accumulated_loss / engine.schedule.num_steps
avg_loss = accumulated_loss / num_steps
lr_scheduler.step()
return avg_loss
@@ -59,25 +64,17 @@ def train(engine):
@pytest.mark.skip("This test should be invoked by test.sh in the same folder as it runs on multiple gpus")
def test_2d_parallel_vision_transformer():
# init dist
model, train_dataloader, test_dataloader, criterion, optimizer, schedule, lr_scheduler = colossalai.initialize(
CONFIG_PATH)
engine, train_dataloader, test_dataloader = colossalai.initialize(CONFIG_PATH)
lr_scheduler = build_lr_scheduler(gpc.config.lr_scheduler, engine.optimizer)
logger = get_global_dist_logger()
engine = Engine(model=model,
train_dataloader=train_dataloader,
test_dataloader=test_dataloader,
criterion=criterion,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
schedule=schedule)
logger.info('start training')
for epoch in range(gpc.config.num_epochs):
train_loss = train(engine)
train_loss = train(engine, train_dataloader, lr_scheduler)
logger.info(f'epoch {epoch} - train loss: {train_loss}')
if epoch % 2 == 0:
correct_sum, total_sum, eval_loss = eval(engine)
correct_sum, total_sum, eval_loss = eval(engine, test_dataloader)
logger.info(
f'epoch {epoch} - eval loss: {eval_loss}, total: {total_sum}, '
f'correct: {correct_sum}, acc: {correct_sum / total_sum}')

42
tests/test_models/test_vision_transformer/test_vit_2p5d/test_vit_2p5d.py
View File

@@ -4,22 +4,25 @@ import pytest
import torch.autograd
import colossalai
from colossalai.builder import build_lr_scheduler
from colossalai.context.parallel_mode import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.engine import Engine
from colossalai.logging import get_global_dist_logger
from colossalai.nn.layer._parallel_utilities import _gather
CONFIG_PATH = Path(__file__).parent.parent.joinpath('configs/vit_2p5d.py')
def eval(engine):
def eval(engine, test_dataloader):
engine.eval()
accumulated_loss = 0
correct_sum = 0
total_sum = 0
num_steps = len(test_dataloader)
data_iter = iter(test_dataloader)
for i in range(engine.schedule.num_steps):
output, label, loss = engine.step()
for i in range(num_steps):
output, label, loss = engine.step(data_iter)
accumulated_loss += loss.detach().cpu().numpy()
output = _gather(
@@ -41,18 +44,21 @@ def eval(engine):
correct = torch.sum(label[0] == output)
correct_sum += correct
total_sum += label[0].size(0)
avg_loss = accumulated_loss / engine.schedule.num_steps
avg_loss = accumulated_loss / num_steps
return correct_sum, total_sum, avg_loss
def train(engine):
def train(engine, train_dataloader, lr_scheduler):
engine.train()
accumulated_loss = 0
num_steps = len(train_dataloader)
data_iter = iter(train_dataloader)
for i in range(engine.schedule.num_steps):
output, label, loss = engine.step()
for i in range(num_steps):
output, label, loss = engine.step(data_iter)
accumulated_loss += loss.detach().cpu().numpy()
avg_loss = accumulated_loss / engine.schedule.num_steps
avg_loss = accumulated_loss / num_steps
lr_scheduler.step()
return avg_loss
@@ -60,29 +66,21 @@ def train(engine):
@pytest.mark.skip("This test should be invoked by test.sh in the same folder as it runs on multiple gpus")
def test_2p5d_parallel_vision_transformer():
# init dist
model, train_dataloader, test_dataloader, criterion, optimizer, schedule, lr_scheduler = colossalai.initialize(
CONFIG_PATH)
engine, train_dataloader, test_dataloader = colossalai.initialize(CONFIG_PATH)
lr_scheduler = build_lr_scheduler(gpc.config.lr_scheduler, engine.optimizer)
logger = get_global_dist_logger()
engine = Engine(model=model,
train_dataloader=train_dataloader,
test_dataloader=test_dataloader,
criterion=criterion,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
schedule=schedule)
logger.info('start training')
for epoch in range(gpc.config.num_epochs):
train_loss = train(engine)
train_loss = train(engine, train_dataloader, lr_scheduler)
logger.info(f'epoch {epoch} - train loss: {train_loss}')
if epoch % 2 == 0:
correct_sum, total_sum, eval_loss = eval(engine)
correct_sum, total_sum, eval_loss = eval(engine, test_dataloader)
logger.info(
f'epoch {epoch} - eval loss: {eval_loss}, total: {total_sum}, '
f'correct: {correct_sum}, acc: {correct_sum / total_sum}')
if __name__ == '__main__':
test_2p5d_parallel_vision_transformer()
test_2p5d_parallel_vision_transformer()

27
tests/test_models/test_vision_transformer/test_vit_3d/test_vit_3d.py
View File

@@ -1,16 +1,14 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import time
from pathlib import Path
import torch
from tqdm import tqdm
from colossalai import initialize
import colossalai
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.engine import Engine
from colossalai.logging import get_global_dist_logger
from colossalai.trainer import Trainer
from colossalai.trainer.metric import Accuracy3D
@@ -29,7 +27,7 @@ def _train_epoch(epoch, engine):
num_samples = 0
now = time.time()
epoch_start = now
progress = range(engine.schedule.num_steps)
progress = range(engine._schedule.num_steps)
if gpc.get_global_rank() == 0:
progress = tqdm(progress, desc='[Epoch %d]' % epoch, miniters=1)
for step in progress:
@@ -68,7 +66,7 @@ def _eval(epoch, engine):
ParallelMode.PARALLEL_3D_WEIGHT)
total = 0
with torch.no_grad():
for _ in range(engine.schedule.num_steps):
for _ in range(engine._schedule.num_steps):
outputs, targets, loss = engine.step()
if isinstance(outputs, (list, tuple)):
outputs = outputs[0]
@@ -80,32 +78,25 @@ def _eval(epoch, engine):
print_rank_0(
'[Epoch %d] Evaluation loss: %.3f | Acc: %.3f%%' %
(epoch, eval_loss / engine.schedule.num_steps,
(epoch, eval_loss / engine._schedule.num_steps,
acc.get_accumulated_value() * 100), logger)
def train():
model, train_dataloader, test_dataloader, criterion, \
optimizer, schedule, lr_scheduler = initialize(CONFIG_PATH)
# init dist
engine, train_dataloader, test_dataloader = colossalai.initialize(CONFIG_PATH)
logger = get_global_dist_logger()
engine = Engine(model=model,
train_dataloader=train_dataloader,
test_dataloader=test_dataloader,
criterion=criterion,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
schedule=schedule)
logger.info("Engine is built", ranks=[0])
trainer = Trainer(engine=engine, hooks_cfg=gpc.config.hooks, verbose=True)
trainer = Trainer(engine=engine, verbose=True)
logger.info("Trainer is built", ranks=[0])
logger.info("Train start", ranks=[0])
trainer.fit(train_dataloader=train_dataloader,
test_dataloader=test_dataloader,
max_epochs=gpc.config.num_epochs,
epochs=gpc.config.num_epochs,
hooks_cfg=gpc.config.hooks,
display_progress=True,
test_interval=1)