[BUG] Memory is enough for training by using zero-3, but OOM occurred after enabling DeepCompile

Describe the bug
I am trying to test the DeepCompile by the script in DeepSpeedExamples repo

My training model is 'meta-llama/Llama-3.2-1B-Instruct'. My machine has 2*15360MiB Tesla T4 GPUs.

As I want to compare zero-3 with deepcompile, I modify the script like this

PROFILE_DIR=${PROFILE_DIR:-"profiles"}
mkdir -p ${PROFILE_DIR}
PROFILE_OPTS="--profile --profile-dir ${PROFILE_DIR}"
COMPILE_OPTS="--compile"
DC_OPTS="--compile --deepcompile --debug-log"
ACC_OPTS="--gradient-accumulation-steps 1"
AC_OPTS="--activation-checkpointing"

export NUM_NODES=${NUM_NODES:-4}

MODEL=${MODEL_NAME:-"meta-llama/Llama-3.2-1B-Instruct"}
BATCH_SIZE_OPTS=(1)
SEQ_LENGTH_OPTS=(512)
for BATCH_SIZE in ${BATCH_SIZE_OPTS[@]}; do
    for SEQ_LENGTH in ${SEQ_LENGTH_OPTS[@]}; do
        ARGS="--model ${MODEL} --batch-size ${BATCH_SIZE} --seq-length ${SEQ_LENGTH} ${ACC_OPTS} ${AC_OPTS} ${PROFILE_OPTS}"
        bash ./run_multinode.sh --backend deepspeed ${ARGS}
        bash ./run_multinode.sh --backend deepspeed ${ARGS} ${DC_OPTS} --passes prefetch
    done
done

Then I run the script using NUM_NODES=1 NGPUS_PER_NODE=2 MODEL_NAME="meta-llama/Llama-3.2-1B-Instruct" bash run_bench.sh. The training without deepcompile success, but OOM occurs after enabling DeepCompile.

FYI, I have selected some of the output results as follows:
Without DeepCompile

Epoch 1, Step 29, Loss: 0.3229782283306122 sync: True time: 3.6913506984710693 alloc_mem: 11038133248 peak_mem: 14176305664
[2025-07-18 17:30:08,948] [WARNING] [stage3.py:2164:step] 4 pytorch allocator cache flushes since last step. this happens when there is high memory pressure and is detrimental to performance. if this is happening frequently consider adjusting settings to reduce memory consumption. If you are unable to make the cache flushes go away consider adding get_accelerator().empty_cache() calls in your training loop to ensure that all ranks flush their caches at the same time
Epoch 1, Step 30, Loss: 0.33652323484420776 sync: True time: 3.6097521781921387 alloc_mem: 11038133248 peak_mem: 14176305664
meta-llama/Llama-3.2-1B-Instruct ds=True np=2 batch_size=1 seq=512 zero_stage=3 acc=1 ac=True compile=False backend=inductor deepcompile=False passes=None compile_time=0 iteration time: 3.6117 alloc_mem: 11038133248 peak_mem: 14176305664

With DeepCompile

Epoch 1, Step 1, Loss: 6.230960845947266 sync: True time: 66.6338963508606 alloc_mem: 13038156800 peak_mem: 13524622336
[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/lirui/DeepSpeedExamples/benchmarks/deepcompile/run_bench_lm.py", line 276, in <module>
[rank1]:     main()
[rank1]:   File "/home/lirui/DeepSpeedExamples/benchmarks/deepcompile/run_bench_lm.py", line 213, in main
[rank1]:     outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids, use_cache=False)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]:     ret_val = func(*args, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank1]:     loss = self.module(*inputs, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1737, in _wrapped_call_impl
[rank1]:     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 574, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/transformers/utils/generic.py", line 953, in wrapper
[rank1]:     @wraps(func)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 745, in _fn
[rank1]:     return fn(*args, **kwargs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/aot_autograd.py", line 1184, in forward
[rank1]:     return compiled_fn(full_args)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 310, in runtime_wrapper
[rank1]:     all_outs = call_func_at_runtime_with_args(
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/utils.py", line 126, in call_func_at_runtime_with_args
[rank1]:     out = normalize_as_list(f(args))
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/utils.py", line 100, in g
[rank1]:     return f(*args)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/autograd/function.py", line 575, in apply
[rank1]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1585, in forward
[rank1]:     fw_outs = call_func_at_runtime_with_args(
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/utils.py", line 126, in call_func_at_runtime_with_args
[rank1]:     out = normalize_as_list(f(args))
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 490, in wrapper
[rank1]:     return compiled_fn(runtime_args)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 672, in inner_fn
[rank1]:     outs = compiled_fn(args)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_inductor/output_code.py", line 466, in __call__
[rank1]:     return self.current_callable(inputs)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_inductor/utils.py", line 2128, in run
[rank1]:     return model(new_inputs)
[rank1]:   File "/tmp/torchinductor_lirui/4w/c4wsbb5movgt7hl36crty2lqccmmpfabyfzpgnmb63ngghv5sa6p.py", line 1559, in call
[rank1]:     buf288 = torch.ops.dc.allgather_param.default(primals_2, 140545491081296, 0)
[rank1]:   File "/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch/_ops.py", line 723, in __call__
[rank1]:     return self._op(*args, **kwargs)
[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 502.00 MiB. GPU 1 has a total capacity of 14.57 GiB of which 312.75 MiB is free. Including non-PyTorch memory, this process has 14.26 GiB memory in use. Of the allocated memory 13.96 GiB is allocated by PyTorch, and 48.74 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

To Reproduce
Steps to reproduce the behavior:

Compare the peak memory usage before and after enabling DeepCompile. When memory is not sufficiently abundant, enabling DeepCompile may lead to OOM.

ds_report output

(deepcompile) lirui@xxx:~/DeepSpeedExamples/benchmarks/deepcompile$ ds_report
[2025-07-18 17:34:47,832] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-07-18 17:34:50,796] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
dc ..................... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
 [WARNING]  NVIDIA Inference is only supported on Ampere and newer architectures
fp_quantizer ........... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.6
 [WARNING]  using untested triton version (3.2.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/torch']
torch version .................... 2.6.0+cu124
deepspeed install path ........... ['/home/lirui/miniconda3/envs/deepcompile/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.17.2, unknown, unknown
torch cuda version ............... 12.4
torch hip version ................ None
nvcc version ..................... 12.4
deepspeed wheel compiled w. ...... torch 2.6, cuda 12.4
shared memory (/dev/shm) size .... 251.88 GB

System info (please complete the following information):

• OS: Ubuntu 20.04.6 LTS
• GPU count and types : 1 machine with 2 Tesla T4, 15360MiB each
• Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
• Python version: 3.9
• Any other relevant info about your setup

Launcher context
the script use the accelerate

Docker context
no

Additional context
my environment's info is

(deepcompile) lirui@xxx:~/DeepSpeedExamples/benchmarks/deepcompile$ pip list
Package                  Version
------------------------ -----------
accelerate               1.8.1
aiohappyeyeballs         2.6.1
aiohttp                  3.12.14
aiosignal                1.4.0
annotated-types          0.7.0
async-timeout            5.0.1
attrs                    25.3.0
certifi                  2025.7.14
charset-normalizer       3.4.2
datasets                 3.1.0
deepspeed                0.17.2
dill                     0.3.8
einops                   0.8.1
filelock                 3.18.0
frozenlist               1.7.0
fsspec                   2024.9.0
hf-xet                   1.1.5
hjson                    3.1.0
huggingface-hub          0.33.4
idna                     3.10
importlib_metadata       8.7.0
Jinja2                   3.1.6
MarkupSafe               3.0.2
mpmath                   1.3.0
msgpack                  1.1.1
multidict                6.6.3
multiprocess             0.70.16
networkx                 3.2.1
ninja                    1.11.1.4
numpy                    2.0.2
nvidia-cublas-cu12       12.4.5.8
nvidia-cuda-cupti-cu12   12.4.127
nvidia-cuda-nvrtc-cu12   12.4.127
nvidia-cuda-runtime-cu12 12.4.127
nvidia-cudnn-cu12        9.1.0.70
nvidia-cufft-cu12        11.2.1.3
nvidia-cufile-cu12       1.11.1.6
nvidia-curand-cu12       10.3.5.147
nvidia-cusolver-cu12     11.6.1.9
nvidia-cusparse-cu12     12.3.1.170
nvidia-cusparselt-cu12   0.6.2
nvidia-ml-py             12.575.51
nvidia-nccl-cu12         2.21.5
nvidia-nvjitlink-cu12    12.4.127
nvidia-nvtx-cu12         12.4.127
packaging                25.0
pandas                   2.3.1
pillow                   11.3.0
pip                      25.1
propcache                0.3.2
psutil                   7.0.0
py-cpuinfo               9.0.0
pyarrow                  20.0.0
pydantic                 2.11.7
pydantic_core            2.33.2
python-dateutil          2.9.0.post0
pytz                     2025.2
PyYAML                   6.0.2
regex                    2024.11.6
requests                 2.32.4
safetensors              0.5.3
setuptools               78.1.1
six                      1.17.0
sympy                    1.13.1
tokenizers               0.21.2
torch                    2.6.0
torchaudio               2.6.0
torchvision              0.21.0
tqdm                     4.67.1
transformers             4.52.4
triton                   3.2.0
typing_extensions        4.14.1
typing-inspection        0.4.1
tzdata                   2025.2
urllib3                  2.5.0
wheel                    0.45.1
xxhash                   3.5.0
yarl                     1.20.1
zipp                     3.23.0