[BUG] RuntimeError: tracing error at step X, due to `NOT_AVAILABLE` parameters

Describe the bug

During evaluation of a Mistral-Small-24B-Instruct-2501 I have code resembling the following:

from contextlib import contextmanager

import torch.distributed as dist
from deepspeed import DeepSpeedEngine
from transformers import PreTrainedModel
from trl.models.utils import unwrap_model_for_generation


@contextmanager
def prep_for_forward(model: PreTrainedModel | DeepSpeedEngine) -> Iterator[None]:
    if isinstance(model, DeepSpeedEngine):
        model = model.module
    model.train()
    # Note that we are both left- (prompt) and right- (completion) padding
    prev_cache_setting = model.config.use_cache
    model.config.use_cache = False
    yield
    model.config.use_cache = prev_cache_setting


model: PreTrainedModel
with unwrap_model_for_generation(model, accelerator) as local_model:
    prompt_completion_ids = local_model.generate(  # (B, P + C)
        **prompt_inputs,
        synced_gpus=dist.is_initialized(),
        generation_config=gen_config,
        use_cache=True,
        output_logits=False,
        ...,
    )
...

with prep_for_forward(model):
    output = model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        ...
    )

Running this code I get the following RuntimeError on all nodes, the during PreTrainedModel.__call__ (end of the above snippet):

2: [rank18]:   File "/path/to/repo/a.py", line 248, in get_per_token_logps
2: [rank18]:     model(
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapp
ed_call_impl
2: [rank18]:     return self._call_impl(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_
impl
2: [rank18]:     return forward_call(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2: [rank18]:     ret_val = func(*args, **kwargs)
2: [rank18]:               ^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2030, in forwa
rd
2: [rank18]:     loss = self.module(*inputs, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapp
ed_call_impl
2: [rank18]:     return self._call_impl(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_
impl
2: [rank18]:     return inner()
2: [rank18]:            ^^^^^^^
0: Traceback (most recent call last):
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
2: [rank18]:     result = forward_call(*args, **kwargs)
2: [rank18]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/transformers/utils/deprecation.py", line 172, in
wrapped_func
2: [rank18]:     return func(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/transformers/models/mistral/modeling_mistral.py",
 line 843, in forward
2: [rank18]:     outputs = self.model(
2: [rank18]:               ^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapp
ed_call_impl
2: [rank18]:     return self._call_impl(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_
impl
2: [rank18]:     return inner()
2: [rank18]:            ^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
2: [rank18]:     result = forward_call(*args, **kwargs)
2: [rank18]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/transformers/models/mistral/modeling_mistral.py",
 line 554, in forward
2: [rank18]:     layer_outputs = self._gradient_checkpointing_func(
2: [rank18]:                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/_compile.py", line 32, in inner
2: [rank18]:     return disable_fn(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 632, in _fn
2: [rank18]:     return fn(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/utils/checkpoint.py", line 496, in checkpoi
nt
2: [rank18]:     ret = function(*args, **kwargs)
2: [rank18]:           ^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapp
ed_call_impl
2: [rank18]:     return self._call_impl(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_
impl
2: [rank18]:     return inner()
2: [rank18]:            ^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
2: [rank18]:     result = forward_call(*args, **kwargs)
2: [rank18]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/transformers/models/mistral/modeling_mistral.py",
 line 263, in forward
2: [rank18]:     hidden_states = self.mlp(hidden_states)
2: [rank18]:                     ^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapp
ed_call_impl
2: [rank18]:     return self._call_impl(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_
impl
2: [rank18]:     return inner()
2: [rank18]:            ^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1790, in inner
2: [rank18]:     result = forward_call(*args, **kwargs)
2: [rank18]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/transformers/models/mistral/modeling_mistral.py",
 line 57, in forward
2: [rank18]:     down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
2: [rank18]:                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapp
ed_call_impl
2: [rank18]:     return self._call_impl(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1844, in _call_
impl
2: [rank18]:     return inner()
2: [rank18]:            ^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1779, in inner
2: [rank18]:     args_result = hook(self, args)
2: [rank18]:                   ^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2: [rank18]:     ret_val = func(*args, **kwargs)
2: [rank18]:               ^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/runtime/zero/parameter_offload.py", lin
e 292, in _pre_forward_module_hook
2: [rank18]:     self.pre_sub_module_forward_function(module)
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorat
e_context
2: [rank18]:     return func(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/runtime/zero/parameter_offload.py", lin
e 467, in pre_sub_module_forward_function
2: [rank18]:     param_coordinator.fetch_sub_module(sub_module, forward=True)
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 632, in _fn
2: [rank18]:     return fn(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
2: [rank18]:     ret_val = func(*args, **kwargs)
2: [rank18]:               ^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorat
e_context
2: [rank18]:     return func(*args, **kwargs)
2: [rank18]:            ^^^^^^^^^^^^^^^^^^^^^
2: [rank18]:   File "/path/to/repo/.venv/lib/python3.12/site-packages/deepspeed/runtime/zero/partitioned_param_coordina
tor.py", line 365, in fetch_sub_module
2: [rank18]:     raise RuntimeError(
2: [rank18]: RuntimeError: tracing error at step 16:
2: [rank18]: module id: 68919, training: True
2: [rank18]: expected the next 6 parameters in the parameter fetch queue to be ({'id': 'name=unknown id=353', 'status': 'AVAILABLE',
'numel': 5242880, 'ds_numel': 5242880, 'shape': (1024, 5120), 'ds_shape': (1024, 5120), 'requires_grad': True, 'grad_shape': None, 'p
ersist': False, 'active_sub_modules': {68919}, 'ds_tensor.shape': torch.Size([218454])}, {'id': 'name=unknown id=352', 'status': 'AVA
ILABLE', 'numel': 20971520, 'ds_numel': 20971520, 'shape': (4096, 5120), 'ds_shape': (4096, 5120), 'requires_grad': True, 'grad_shape
': None, 'persist': False, 'active_sub_modules': {68919}, 'ds_tensor.shape': torch.Size([873814])}, {'id': 'name=unknown id=7', 'stat
us': 'AVAILABLE', 'numel': 167772160, 'ds_numel': 167772160, 'shape': (5120, 32768), 'ds_shape': (5120, 32768), 'requires_grad': True
, 'grad_shape': None, 'persist': False, 'active_sub_modules': {68919}, 'ds_tensor.shape': torch.Size([6990507])}, {'id': 'name=unknow
n id=354', 'status': 'AVAILABLE', 'numel': 5242880, 'ds_numel': 5242880, 'shape': (1024, 5120),
2: 'ds_shape': (1024, 5120), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {68919}, 'ds_tensor.s
hape': torch.Size([218454])}, {'id': 'name=unknown id=355', 'status': 'AVAILABLE', 'numel': 20971520, 'ds_numel': 20971520, 'shape':
(5120, 4096), 'ds_shape': (5120, 4096), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {68919}, '
ds_tensor.shape': torch.Size([873814])}, {'id': 'name=unknown id=356', 'status': 'AVAILABLE', 'numel': 167772160, 'ds_numel': 1677721
60, 'shape': (32768, 5120), 'ds_shape': (32768, 5120), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modul
es': {68919}, 'ds_tensor.shape': torch.Size([6990507])})
2: [rank18]: but got
2: [rank18]:  ({'id': 'name=unknown id=12', 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 5242880, 'shape': (0,), 'ds_shape': (1
024, 5120), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': set(), 'ds_tensor.shape': torch.Size([
218454])}, {'id': 'name=unknown id=10', 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 20971520, 'shape': (0,), 'ds_shape': (4096
, 5120), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': set(), 'ds_tensor.shape': torch.Size([873
814])}, {'id': 'name=unknown id=17', 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 5120, 'shape': (0,), 'ds_shape': (5120,), 're
quires_grad': True, 'grad_shape': None, 'persist': True, 'active_sub_modules': set(), 'ds_tensor.shape': torch.Size([214])}, {'id': '
name=unknown id=11', 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 5242880, 'shape': (0,), 'ds_shape': (1024, 5120), 'requires_g
rad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': set(), 'ds_tensor.shape'
2: : torch.Size([218454])}, {'id': 'name=unknown id=13', 'status': 'NOT_AVAILABLE', 'numel': 0, 'ds_numel': 20971520, 'shape': (0,),
'ds_shape': (5120, 4096), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': set(), 'ds_tensor.shape'
: torch.Size([873814])}, {'id': 'name=unknown id=7', 'status': 'AVAILABLE', 'numel': 167772160, 'ds_numel': 167772160, 'shape': (5120
, 32768), 'ds_shape': (5120, 32768), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {68919}, 'ds_
tensor.shape': torch.Size([6990507])}).

Looking at this final error message, we see many of the parameters are shown as NOT_AVAILABLE, and I am not sure why.

To Reproduce

Sorry this is not reproducible for now, if needed I can try to make a repro.

Please just go based on concepts if possible.

Expected behavior

The sharded model.__call__ to have available parameters after unwrap_model_for_generation's context completes, and be able to recreate the logprobs.

ds_report output

--------------------------------------------------
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]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
 [WARNING]  FP Quantizer is using an untested triton version (3.1.0), only 2.3.(0, 1) and 3.0.0 are known to be compatible with these kernels
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.5
 [WARNING]  using untested triton version (3.1.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]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/path/to/.venv/lib/python3.12/site-packages/torch']
torch version .................... 2.5.1+cu124
deepspeed install path ........... ['/path/to/.venv/lib/python3.12/site-packages/deepspeed']
deepspeed info ................... 0.16.5, unknown, unknown
torch cuda version ............... 12.4
torch hip version ................ None
nvcc version ..................... 12.4
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
shared memory (/dev/shm) size .... 503.69 GB

System info (please complete the following information):

• OS: Ubuntu 22.04 LTS
• GPU count and types: three nodes with 8 H100s each
• Hugging Face Transformers/Accelerate/etc. versions: accelerate==1.4.0, transformers==4.49.0, trl==0.16.0
• Python version: 3.12.9
• Any other relevant info about your setup

Additional context

Happy to provide more if needed.