github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-11 22:10:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

[misc] Accelerate CI for zero and dist optim (#5758)

Browse Source 
* remove fp16 from lamb

* remove d2h copy in checking states

---------

Co-authored-by: Edenzzzz <wtan45@wisc.edu>
This commit is contained in:
Edenzzzz
2024-06-05 11:25:19 +08:00
committed by GitHub
parent 50b4c8e8cf
commit 79f7a7b211
12 changed files with 65 additions and 400 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_optimizer/test_dist_lamb.py
View File

@@ -18,7 +18,6 @@ from tests.test_optimizer._utils import check_optim_states, run_bert_test
_ALLOWED_P_G_TYPES = [
(torch.float, torch.float), # pure fp32
(torch.float, torch.half), # fp16 amp
(torch.float, torch.bfloat16), # bfloat16 amp
]
@@ -264,7 +263,6 @@ def run_dist_lamb_fwd_bwd(
torch_optim.step()
optim.step()
dist.barrier()
torch_optim.zero_grad()
optim.zero_grad()
try: