[BUG] DeepCompile: Training hang on random-sized inputs

[eternalNight]

Description

Real-life training data may not be of the same size for every rank and at every iteration. When DeepCompile is active, training with variable-length data can hang because DeepCompile requires communication among the ranks during profiling, but:

1. The compute graph may not be exactly the same across ranks (e.g. some have specific padding while the others don't).
2. Guard failure (due to tensor size change) on different ranks may occur at different iterations.

To Reproduce

1. Download https://gist.github.com/eternalNight/3c2cf8c703f1e9e7742d3b7f9e1edae3
2. Execute deepspeed --num_gpus=N openvla-like.py -c -r.

Expected behavior

DeepCompile works for variable-length training data.

[Shreyas-S-809]

Hey I would like to take up this issue, Could you assign me ?

[eternalNight]

Hey I would like to take up this issue, Could you assign me ?

Sure, and I appreciate your help! This is fundamental to DeepCompile which needs to reduce profiling data to ensure all ranks share the same order. Feel free to share your insights for further discussions.

[Shreyas-S-809]

Hey I would like to take up this issue, Could you assign me ?

Sure, and I appreciate your help! This is fundamental to DeepCompile which needs to reduce profiling data to ensure all ranks share the same order. Feel free to share your insights for further discussions.

Thanks for assigning this!

I’ve started digging into the issue, and my current understanding is that the hang is likely caused by cross-rank graph divergence when using variable-length inputs — leading to inconsistent guard failures and mismatched profiling order across ranks.

I’ll spend some time reproducing this locally and validating the exact failure point (graph differences vs. guard-triggered recompilation vs. profiling sync).

I’ll share a detailed update along with a proposed fix once I have a concrete direction.

---

I had a couple of clarifying questions:

1. Is DeepCompile expected to fully support variable-length inputs across ranks, or is some level of shape consistency (e.g., padding/bucketing) currently assumed?

2. During profiling, does DeepCompile strictly require identical execution order across ranks, or is there any existing mechanism to tolerate minor divergence?

3. Are there any known design discussions or prior attempts around synchronizing recompilation / guard failures across ranks?

---

Happy to iterate based on your guidance — will follow up soon with findings.

[eternalNight]

I had a couple of clarifying questions:

1. Is DeepCompile expected to fully support variable-length inputs across ranks, or is some level of shape consistency (e.g., padding/bucketing) currently assumed?

DeepCompile does not support variable-length inputs today, but it is definitely important to add that support as real-world datasets are barely fixed-sized.

The key issue here is that dynamo has no idea which dimension in which input tensor is variable, and thus it has to trace a graph with fixed sizes for the first time and later re-trace it when a different dimension is seen. During distributed training different ranks may seen the variable dimension (and recompile the graph) at different iterations. But DeepCompile today is designed to allreduce the profiled execution time and memory usage among all ranks. With only part of the ranks recompiling the graph, such allreduce will block the training process.

2. During profiling, does DeepCompile strictly require identical execution order across ranks, or is there any existing mechanism to tolerate minor divergence?

The key is to make sure collective operations on the same cuda stream or NCCL communicator always have the same order and grouped in the same way. You can check whether the op scheduler (see deepspeed/compile/list_schedule.py) satisfies that property.

3. Are there any known design discussions or prior attempts around synchronizing recompilation / guard failures across ranks?

I'm not aware of any. @tohtana any ideas?

Happy to iterate based on your guidance — will follow up soon with findings.

[tohtana]

Hi @Shreyas-S-809, @eternalNight,

Thank you for working on this issue!
I actually tried to support variable-length inputs and finally ended up padding. Here are my findings:

• Padding output must have matching the same tensor layout, not only tensor shape (The usage in OpenRLHF)
• Even after padding to the maximum length across multiple ranks, the length changes at every iteration. Dynamo finds a guard failure and relaxes the guard. At certain point, it will reach a stable graph.
• If I remember correctly, Dynamo had an (undocument) option to exchange information across ranks. It seemed to be designed to get stable compiled graphs across ranks, but it didn't work for me. As Dynamo has significantly updated since then, we should search for such an option again.

I hope this experience useful to resolve this issue. I'm also happy to work together if I have anything I could do.

[Shreyas-S-809]

Thanks for the detailed insights @eternalNight , @tohtana — this clarifies the problem significantly.

From your explanations, my understanding is that the core issue is desynchronized recompilation across ranks due to variable input dimensions, which leads to mismatched profiling states and ultimately breaks collective operation ordering.

I’ll start by:

• Reproducing the issue and validating where recompilation diverges across ranks

• Inspecting the op scheduler to verify collective ordering guarantees

• Exploring ways to enforce synchronized compilation or consistent graph states across ranks

This might take some time as I want to ensure the solution aligns with DeepCompile’s design constraints, but I’ll share progress and findings as I move forward.

Really appreciate the guidance — this is very helpful.

[Shreyas-S-809]

Hi @Shreyas-S-809, @eternalNight,

Thank you for working on this issue! I actually tried to support variable-length inputs and finally ended up padding. Here are my findings:

• Padding output must have matching the same tensor layout, not only tensor shape (The usage in OpenRLHF)
• Even after padding to the maximum length across multiple ranks, the length changes at every iteration. Dynamo finds a guard failure and relaxes the guard. At certain point, it will reach a stable graph.
• If I remember correctly, Dynamo had an (undocument) option to exchange information across ranks. It seemed to be designed to get stable compiled graphs across ranks, but it didn't work for me. As Dynamo has significantly updated since then, we should search for such an option again.

I hope this experience useful to resolve this issue. I'm also happy to work together if I have anything I could do.

Hi @eternalNight , @tohtana , I was able to reproduce a failure when enabling DeepCompile with ZeRO stage 1.

With the following config:

"zero_optimization": { "stage": 1 },
"compile": { "deepcompile": true }

Running:
deepspeed repro.py -c -r -b 1 -i 5

Results in:

ValueError: A pass for ZeRO3 is not specified though ZeRO1 is enabled

From debugging, this originates in init_z1, which appears to expect ZeRO3-specific passes even when ZeRO1 is enabled.

Could you confirm:

1. Is DeepCompile expected to support ZeRO stage 1?
2. Or should this configuration be disallowed with a clearer error message?

take a look at the OP :

https://docs.google.com/document/d/18NO7A-P5ExL0vTxwSAzn6HixZRQ386UP-zBdj24tAdg/edit?usp=sharing

Notebook :

https://colab.research.google.com/drive/1PSs4KrqRogXkUi11BMdCEGbjE7sdMNMB?usp=sharing

Happy to submit a PR depending on the intended behavior.

[eternalNight]

From debugging, this originates in init_z1, which appears to expect ZeRO3-specific passes even when ZeRO1 is enabled.

Could you confirm:

1. Is DeepCompile expected to support ZeRO stage 1?

DeepCompile does support ZeRO1. It uses a different pass for ZeRO1/2 and ZeRO3 (zero1_compile and zero3_compile, respectively). So if you explicitly pass a compile schedule, you should make sure the inclusion of zero1_compile or zero3_compile is consistent with your ZeRO setting.

[Shreyas-S-809]

From debugging, this originates in init_z1, which appears to expect ZeRO3-specific passes even when ZeRO1 is enabled.
Could you confirm:

1. Is DeepCompile expected to support ZeRO stage 1?

DeepCompile does support ZeRO1. It uses a different pass for ZeRO1/2 and ZeRO3 (zero1_compile and zero3_compile, respectively). So if you explicitly pass a compile schedule, you should make sure the inclusion of zero1_compile or zero3_compile is consistent with your ZeRO setting.

Thanks for the clarification!

I’m not explicitly passing a compile schedule in my repro, so it seems the default schedule may be including a ZeRO3-specific pass when running with ZeRO1.

I’ll investigate how the schedule is being constructed and verify whether zero3_compile is being included unintentionally.

Will follow up with findings or a PR if I identify the root cause.

[eternalNight]

I’m not explicitly passing a compile schedule in my repro, so it seems the default schedule may be including a ZeRO3-specific pass when running with ZeRO1.

From the calltrace, I see model.compile(backend=backend, schedule=schedule) on line 290 in your repro.py. That schedule=schedule is the compile schedule that the engine uses.

[Shreyas-S-809]

I’m not explicitly passing a compile schedule in my repro, so it seems the default schedule may be including a ZeRO3-specific pass when running with ZeRO1.

From the calltrace, I see model.compile(backend=backend, schedule=schedule) on line 290 in your repro.py. That schedule=schedule is the compile schedule that the engine uses.

Thanks for pointing that out — you're right.

I was passing a schedule in my repro, and it included a ZeRO3-specific pass while running with ZeRO1, which caused the issue.

I’ll correct the schedule to use zero1_compile (or remove it entirely to rely on defaults) and verify.

Appreciate the clarification!

[Shreyas-S-809]

I’m not explicitly passing a compile schedule in my repro, so it seems the default schedule may be including a ZeRO3-specific pass when running with ZeRO1.

From the calltrace, I see model.compile(backend=backend, schedule=schedule) on line 290 in your repro.py. That schedule=schedule is the compile schedule that the engine uses.

DeepCompile + ZeRO1: KeyError during backward graph compilation

Thanks for the earlier clarification — that helped resolve the configuration issue.

I updated the repro to use a ZeRO1-compatible compile setup, and while the config issue is now fixed, I’m encountering a new failure during the backward pass when DeepCompile is enabled.

---

Configuration / Updated Behaviour

• ZeRO stage: 1
• Compile passes: No explicit ZeRO3 passes
• Schedule: Default ZeRO1 compile schedule

---

Error

KeyError: <frame_id>
  File "deepspeed/compile/backend.py", line 350, in make_bw_graph
    frames_needing_bwd.remove(frame_id)

---

Observed Behavior

• Failure occurs during backward graph compilation in the inductor backend.
• The error suggests that frame_id is being removed from frames_needing_bwd when it is not present.

---

Initial Analysis

This appears to indicate a mismatch in how backward frames are tracked. Possible causes:

• Duplicate removal of the same frame_id
• Missing or skipped registration of a frame
• Ordering/synchronization issue during backward graph construction

---

Question

Is this expected behavior with the current DeepCompile + ZeRO1 setup, or does this indicate a bug?

I’m happy to investigate further in:

deepspeed/compile/backend.py::make_bw_graph

and propose a fix if this is unintended.

---

Additional Notes

• Issue reproduced after removing ZeRO3-specific compile passes
• Using default ZeRO1 scheduling path

---

Let me know if any additional logs or a minimal repro would help.