github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-07 11:03:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

[example] gpt demo more accuracy tflops (#2178)

Browse Source
This commit is contained in:
Jiarui Fang 2022-12-22 20:51:35 +08:00 committed by GitHub
parent ab54fed292
commit 65f56f49e8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
examples/language/gpt/train_gpt_demo.py
View File

@ -283,6 +283,7 @@ def main():
optimizer.sync_grad() optimizer.sync_grad()
optimizer.step() optimizer.step()
logger.info(get_mem_info(prefix=f'[{n+1}/{NUM_STEPS}] Optimizer step '), ranks=[0]) logger.info(get_mem_info(prefix=f'[{n+1}/{NUM_STEPS}] Optimizer step '), ranks=[0])
torch.cuda.synchronize()
step_time = time() - start step_time = time() - start
logger.info( logger.info(
f'[{n+1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}', f'[{n+1}/{NUM_STEPS}] Loss:{loss.item():.3f}, Step time: {step_time:.3f}s, TFLOPS: {get_tflops_func(step_time):.3f}',