use strict=false when loading checkpoints with mfsdp (#1390)

Should fix failing TOT recipe tests

Signed-off-by: Peter St. John <[email protected]>

Author: pstjohn

bionemo-recipes/recipes/esm2_native_te/checkpoint.py

@@ -240,7 +240,7 @@ def load_checkpoint_mfsdp(
     }
     torch.distributed.checkpoint.load(state_dict=ckpt_state_dict, checkpoint_id=checkpoint_path)
 
-    model.load_state_dict(ckpt_state_dict["model"])
+    model.load_state_dict(ckpt_state_dict["model"], strict=False)
     optimizer.load_state_dict(ckpt_state_dict["optimizer"])
     scheduler.load_state_dict(ckpt_state_dict["scheduler"])
     dataloader = load_dataloader(dataloader, checkpoint_path, dist_config)

bionemo-recipes/recipes/vit/README.md

@@ -83,7 +83,7 @@ def load_dcp_checkpoint(checkpoint_path, model=None, optimizer=None):
         state_dict["optimizer"] = optimizer.state_dict()
     torch.distributed.checkpoint.load(state_dict, checkpoint_id=checkpoint_path)
     if model is not None:
-        model.load_state_dict(state_dict["model"])
+        model.load_state_dict(state_dict["model"], strict=False)
     if optimizer is not None:
         optimizer.load_state_dict(state_dict["optimizer"])
 ```

bionemo-recipes/recipes/vit/checkpoint.py

@@ -67,7 +67,7 @@ def load_dcp_checkpoint(checkpoint_path, model=None, optimizer=None):
         state_dict["optimizer"] = optimizer.state_dict()
     torch.distributed.checkpoint.load(state_dict, checkpoint_id=checkpoint_path)
     if model is not None:
-        model.load_state_dict(state_dict["model"])
+        model.load_state_dict(state_dict["model"], strict=False)
     if optimizer is not None:
         optimizer.load_state_dict(state_dict["optimizer"])