github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 17:46:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

polish pp middleware (#2476)

Browse Source 
Co-authored-by: Ziyue Jiang <ziyue.jiang@gmail.com>
This commit is contained in:
Ziyue Jiang
2023-01-13 16:56:01 +08:00
committed by GitHub
parent a5dc4253c6
commit fef5c949c3
3 changed files with 3 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
examples/language/gpt/experiments/pipeline_parallel/train_gpt_pp.py
View File

@@ -120,7 +120,7 @@ def run_master(args):
logger.info(f'{rank=} numel in the partition:{numel}')
# build optim
pp_engine.initialize_optimizer(HybridAdam, lr=1e-3)
pp_engine.initialize_optimizer(torch.optim.Adam, lr=1e-3)
ranks_tflops = {}
for n in range(NUM_STEPS):