Thanks @XAheli for digging into this! The root-cause analysis is correct, and keeping ZeROOrderedDict as an eager fallback guarded by torch.compiler.is_compiling() is the right direction. But I think we still need some rework.

One issue is the post-forward release loop just moves the failure into backward (this confirms the earlier P1 review comment). With DeepCompile actually enabled and the parameters not persistent (stage3_param_persistence_threshold: 0), training still crashes in engine.backward().

RuntimeError: The size of tensor a (0) must match the size of tensor b (384) at non-singleton dimension 1

Autograd saves leaf parameters by reference and reads param.data at backward time; ZeRO-3's free_param swaps .data to an empty tensor. In eager ZeRO-3 the pre-backward hooks swap the data back, but DeepCompile removed those hooks, so partitioning every AVAILABLE non-persistent param right after forward frees exactly the weights the skipped frame's eager autograd nodes still need.

Another issue is that this PR releases gathered parameters based on global ds_status state that other components also manage. For example, selective_gather keeps chosen parameters gathered via the C++-side registry without updating Python's ds_persist, so the sweep's not p.ds_persist guard cannot exclude them. This kind of interference is reproducible: with DeepCompile actually enabled, the regression test fails on this branch with KeyError: wait_allgather_ds_param__arg0_1_0 (master fails with the expected 'weight' must be 2-D), and the failure disappears once the sweep is replaced by tracked-set release.

Your tests didn't catch these issues because test_deepcompile_skipped_frame.py never activates DeepCompile (ds_config_z3.json doesn't set "compile": {"deepcompile": true}). Also, all its parameters are below the default persistence threshold (100000 elements). So the release loop excludes them even when DeepCompile is on. The test passes on master without your fix.

To make this concrete, I've opened a PR against your branch implementing the rework: XAheli#1. In summary it:

1. records the params the fallback actually gathers and releases only that set after their eager backward use completes (the all-parameter post-forward sweep is removed);
2. skips register_external_parameter in the fallback when DeepCompile is active (its consumers are removed with the module hooks, so the registration is dead state — eager-mode behavior is unchanged);
3. moves the pre/post logic out of DeepSpeedEngine.forward into DeepCompile-owned preprocess/postprocess in deepspeed/compile/;
4. makes the regression test real: DeepCompile enabled, persistence defeated, asserts the frame was actually skipped, and fails on master without the fix.

On @eternalNight's question: This still increases the peak memory, though the cost is now bounded to the fallback-gathered set. But I think it would avoid errors and keep correctness.

Can you take a look at the rework, adjust it as needed, and merge it into your branch if it looks reasonable? (if you'd prefer to address these issues in your own way, that works just as well) Either way, once the backward-safe release Thanks again for working on this!

@tohtana @eternalNight thanks a lot for the detailed review :) I'll take a deeper look and push the changes soon!

Hi @XAheli, just following up here. I’m thinking of using the repair PR I opened against your branch to unblock this:
XAheli#1

If you’d like to make the changes yourself or prefer a different approach, please let me know. Otherwise I’ll proceed with this repair path so we can complete this PR.

Hello @tohtana ! Apologies for the delay, I'll be pushing the changes by eod.

@tohtana Thanks again for the review and the repair PR :)

I've reworked the fix to incorporate your approach. I confirmed that post-forward partition corrupts backward. param.data = shard between forward and backward breaks autograd because it saves leaf parameters by reference and reads param.data at backward time. Simple repro:

p = torch.nn.Parameter(torch.randn(4, 4))
x = torch.randn(4, 4, requires_grad=True)
y = F.linear(x, p)
p.data = torch.zeros(2)  # simulate partition
y.sum().backward()  # RuntimeError: size mismatch

The fix now releases gathered params in a post-backward hook (deepcompile_backward_epilogue) instead of after forward. Only the params the fallback actually gathered are tracked and released. There's no global ds_status sweep.

register_external_parameter is skipped in DeepCompile mode, since the module hooks that consume _external_params are removed by init_z3().

Tests now activates DeepCompile with the new config (ds_config_z3_deepcompile_no_persist.json) sets "deepcompile": true with stage3_param_persistence_threshold: 0. The model uses vocab_size=384, hidden=384 so all params exceed the 100k threshold. The test also asserts fallback_stats.total_gathered_params > 0 to confirm the fallback path was exercised. Uses @torch.compiler.disable instead of graph_break() in a loop for deterministic behavior. Validated all on 2× H200

• Regression test (DeepCompile ON, persistence=0, skipped frame): PASS
• Existing test_compile.py (100 steps): PASS
• LLaMA + DeepCompile: PASS
• LLaMA with tied embeddings + DeepCompile: PASS

One thing I want to flag is that Qwen2 MoE + DeepCompile + persistence_threshold=0 hits a KeyError: wait_allgather_ds_param__arg0_1_0 in list_schedule.py:get_last_uses(). I tested this with current approach (no _original_parameters restoration) and it's the same error, so it's not from this PR. On upstream without the fix it's masked by the earlier 'weight' must be 2-D crash at the embedding lookup. To me looks like a separate graph scheduling issue with MoE model structure when all persistence is disabled. Happy to dig into that if useful!