fix optimizer resuming in FSDP and remove param/buffer precision

Author: mehdidc

src/training/main.py

@@ -87,7 +87,6 @@ def get_latest_checkpoint(path: str, remote : bool):
 
 def main(args):
     args = parse_args(args)
-
     if torch.cuda.is_available():
         # This enables tf32 on Ampere GPUs which is only 8% slower than
         # float16 and almost as accurate as float32
@@ -306,9 +305,7 @@ def main(args):
                 "fp32": torch.float32,
             }
             mixed_precision = MixedPrecision(
-                param_dtype=type_name_to_class[args.precision],
                 reduce_dtype=type_name_to_class[args.fsdp_gradient_reduction_precision],
-                buffer_dtype=type_name_to_class[args.fsdp_buffer_precision],
             )
             layers = set()
             for module in model.modules():
@@ -408,6 +405,13 @@ def _param_name_without_fsdp_prefix(n):
             scaler = GradScaler() if args.precision == "amp" else None
     # optionally resume from a checkpoint
     start_epoch = 0
+    if args.fsdp:
+        FSDP.set_state_dict_type(
+            model,
+            StateDictType.FULL_STATE_DICT,
+            FullStateDictConfig(rank0_only=False, offload_to_cpu=True),
+                FullOptimStateDictConfig(rank0_only=False, offload_to_cpu=True),
+        )
     if args.resume is not None:
         checkpoint = pt_load(args.resume, map_location='cpu')
         if 'epoch' in checkpoint:
@@ -419,7 +423,10 @@ def _param_name_without_fsdp_prefix(n):
             model.load_state_dict(sd)
             if optimizer is not None:
                 if args.fsdp:
-                    sharded_state_dict = FSDP.optim_state_dict_to_load(checkpoint["optimizer"], model, optimizer)
+                    optimizer_state_dict = checkpoint["optimizer"]
+                    optimizer_state_dict['state']['logit_scale']['exp_avg'] = optimizer_state_dict['state']['logit_scale']['exp_avg'].view(1)
+                    optimizer_state_dict['state']['logit_scale']['exp_avg_sq'] = optimizer_state_dict['state']['logit_scale']['exp_avg_sq'].view(1)
+                    sharded_state_dict = FSDP.optim_state_dict_to_load(model, optimizer, optimizer_state_dict)
                     optimizer.load_state_dict(sharded_state_dict)
                 else:
                     optimizer.load_state_dict(checkpoint["optimizer"])
@@ -431,13 +438,7 @@ def _param_name_without_fsdp_prefix(n):
             model.load_state_dict(checkpoint)
             logging.info(f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})")
 
-    if args.fsdp:
-        FSDP.set_state_dict_type(
-            model,
-            StateDictType.FULL_STATE_DICT,
-            FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
-            FullOptimStateDictConfig(rank0_only=True, offload_to_cpu=True),
-        )
+  
     # initialize datasets
     data = get_data(args, (preprocess_train, preprocess_val), epoch=start_epoch, tokenizer=get_tokenizer(args.model))
     assert len(data), 'At least one train or eval dataset must be specified.'
@@ -506,14 +507,15 @@ def _param_name_without_fsdp_prefix(n):
 
         # Saving checkpoints.
         if args.save_logs:
-            
+
             checkpoint_dict = {
                 "epoch": completed_epoch,
                 "name": args.name,
                 "state_dict": model.state_dict(),
-                "optimizer": FSDP.optim_state_dict(model, optimizer) if args.fsdp else optimizer.state_dict()
+                "optimizer": FSDP.optim_state_dict(model, optimizer) if args.fsdp else optimizer.state_dict(),
 
             }
+
             if scaler is not None:
                 checkpoint_dict["scaler"] = scaler.state_dict()
             if is_master(args):

src/training/params.py

@@ -398,12 +398,6 @@ def parse_args(args):
         nargs='+',
         help="Module names to wrap for gradient checkpointing when FSDP is used",
     )
-    parser.add_argument(
-        "--fsdp-buffer-precision",
-        choices=["bf16", "fp16", "fp32"],
-        default="fp32",
-        help="FSDP floating point precision for buffers"
-    )
     parser.add_argument(
         "--fsdp-gradient-reduction-precision",
         choices=["bf16", "fp16", "fp32"],