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 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>

tests/test_trainer/configs/test_trainer_resnet.py

@@ -3,6 +3,7 @@ from pathlib import Path
 
 BATCH_SIZE = 128
 IMG_SIZE = 32
+num_epochs = 200
 
 # resnet 50
 model = dict(
@@ -77,18 +78,14 @@ hooks = [
     dict(type='AccuracyHook'),
     dict(type='LossHook'),
     dict(type='TensorboardHook', log_dir='./tfb_logs'),
+    dict(
+        type='LRSchedulerHook',
+        by_epoch=True,
+        lr_scheduler_cfg=dict(
+            type='CosineAnnealingLR',
+            warmup_steps=5
+        )
+    ),
     dict(type='SaveCheckpointHook', interval=5, checkpoint_dir='./ckpt'),
-    # dict(type='LoadCheckpointHook', epoch=20, checkpoint_dir='./ckpt')
 ]
 
-# fp16 = dict(
-#     mode=AMP_TYPE.PARALLEL,
-#     initial_scale=1
-# )
-
-lr_scheduler = dict(
-    type='CosineAnnealingLR',
-    T_max=200
-)
-
-num_epochs = 200

tests/test_trainer/configs/test_trainer_vit_2d.py

@@ -11,6 +11,7 @@ NUM_ATTENTION_HEADS = 8
 SUMMA_DIM = 2
 NUM_CLASSES = 10
 DEPTH = 6
+num_epochs = 60
 
 train_data = dict(
     dataset=dict(type='CIFAR10Dataset',
@@ -52,13 +53,6 @@ optimizer = dict(type='Adam', lr=0.001, weight_decay=0)
 
 loss = dict(type='CrossEntropyLoss2D', )
 
-# model = dict(
-#     type='VanillaResNet',
-#     block_type='ResNetBasicBlock',
-#     layers=[2, 2, 2, 2],
-#     num_cls=10
-# )
-
 model = dict(
     type='VisionTransformerFromConfig',
     tensor_splitting_cfg=dict(type='ViTInputSplitter2D', ),
@@ -114,8 +108,15 @@ hooks = [
     dict(type='Accuracy2DHook'),
     dict(type='LossHook'),
     dict(type='TensorboardHook', log_dir='./tfb_logs'),
+    dict(
+        type='LRSchedulerHook',
+        by_epoch=True,
+        lr_scheduler_cfg=dict(
+            type='LinearWarmupLR',
+            warmup_steps=5
+        )
+    ),
     dict(type='SaveCheckpointHook', interval=5, checkpoint_dir='./ckpt'),
-    # dict(type='LoadCheckpointHook', epoch=20, checkpoint_dir='./ckpt')
 ]
 
 parallel = dict(
@@ -125,11 +126,8 @@ parallel = dict(
 
 fp16 = dict(mode=AMP_TYPE.PARALLEL, initial_scale=2 ** 8)
 
-lr_scheduler = dict(type='LinearWarmupLR', warmup_epochs=5)
-
-schedule = dict(num_microbatches=1)
-
-num_epochs = 60
-num_microbatches = 1
+engine = dict(
+    schedule=dict(num_microbatches=1)
+)
 
 logging = dict(root_path='./logs')

tests/test_trainer/test_trainer.py

@@ -1,25 +1,16 @@
 import colossalai
 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
 
 
 def test_trainer():
-    model, train_dataloader, test_dataloader, criterion, optimizer, schedule, lr_scheduler = colossalai.initialize()
+    engine, train_dataloader, test_dataloader = colossalai.initialize()
     logger = get_global_dist_logger()
 
-    engine = Engine(
-        model=model,
-        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)
     logger.info("trainer is built", ranks=[0])
 
@@ -27,7 +18,8 @@ def test_trainer():
     trainer.fit(
         train_dataloader=train_dataloader,
         test_dataloader=test_dataloader,
-        max_epochs=gpc.config.num_epochs,
+        hooks_cfg=gpc.config.hooks,
+        epochs=gpc.config.num_epochs,
         display_progress=False,
         test_interval=5
     )