From 7549f84d1ab2456ec719b42ab95c2eadb8c84eb6 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Wed, 9 Sep 2026 14:11:40 +0200 Subject: [PATCH] Fix hybrid extra state size mismatch Signed-off-by: Evgeny --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index fc482ce5b2..360d4f6b52 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/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():