Fix hybrid extra state size mismatch (#3499)

Signed-off-by: Evgeny <etsykunov@nvidia.com>

Author: negvet

tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py

@@ -2192,8 +2192,13 @@ def run_continuation_step(current_model, current_optimizer):
             with te.autocast(enabled=True, recipe=hybrid_recipe):
                 ref_output = model(x).clone()
 
+        # CustomRecipe's _extra_state pickle size isn't guaranteed across models/
+        # ranks, so strip it like the DelayedScaling case above (#1860).
+        model_state = {
+            k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state")
+        }
         save_state = {
-            "model": model.state_dict(),
+            "model": model_state,
             "optimizer": optimizer.state_dict(),
         }
         dcp.save(save_state, checkpoint_id=checkpoint_dir)
@@ -2213,12 +2218,15 @@ def run_continuation_step(current_model, current_optimizer):
         F.mse_loss(out_tmp, target).backward()
         optimizer2.step()
 
+        model2_state = {
+            k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state")
+        }
         state_to_load = {
-            "model": model2.state_dict(),
+            "model": model2_state,
             "optimizer": optimizer2.state_dict(),
         }
         dcp.load(state_to_load, checkpoint_id=checkpoint_dir)
-        model2.load_state_dict(state_to_load["model"])
+        model2.load_state_dict(state_to_load["model"], strict=False)
         optimizer2.load_state_dict(state_to_load["optimizer"])
 
         with torch.no_grad():