perf: zero_grad(set_to_none=True) and reduce checkpoint I/O (#83)

* perf: use zero_grad(set_to_none=True) to reduce memory bandwidth

* perf: only save HuggingFace checkpoint at final step to reduce I/O stalls

Author: Luodian

training/train.py

@@ -664,7 +664,7 @@ def wrap_ddp(model):
             for pfc in list_module_pfc:
                 clip_grad_norm_(pfc.parameters(), max_norm=5, norm_type=2)
             opt.step()
-            opt.zero_grad()
+            opt.zero_grad(set_to_none=True)
 
         lr_scheduler.step()
 
@@ -692,8 +692,6 @@ def wrap_ddp(model):
                 list_head_names=args.list_head_names,
                 keep_num=20,
             )
-            # Also save in HuggingFace format
-            save_hf_checkpoint(args.output, backbone, global_step=global_step, image_size=args.image_size[0])
 
         if global_step > args.total_steps:
             save_checkpoint(