Skip to content

fix(engine): resolve inference workspace attribute lookup - #8288

Open
nathon-lee wants to merge 1 commit into
deepspeedai:masterfrom
nathon-lee:fix/hybird-engine-release-cache-workspace
Open

fix(engine): resolve inference workspace attribute lookup#8288
nathon-lee wants to merge 1 commit into
deepspeedai:masterfrom
nathon-lee:fix/hybird-engine-release-cache-workspace

Conversation

@nathon-lee

@nathon-lee nathon-lee commented Aug 22, 2026

Copy link
Copy Markdown
Contributor

Description

Fix recursive attribute lookup in DeepSpeedEngine.__getattr__.

WorkspaceOp inherits from torch.nn.Module, so assigning it to
DeepSpeedHybridEngine.workspace registers it in the engine's _modules
mapping. The existing DeepSpeedEngine.__getattr__ implementation did not
delegate to torch.nn.Module.__getattr__ before forwarding missing attributes
to the wrapped model.

When HybridEngine was configured with:

"hybrid_engine": {
    "enabled": True,
    "release_inference_cache": True,
}

the first call to retake_inference_cache() attempted to access
self.workspace. The registered submodule was not resolved correctly, and the
following logic recursively invoked DeepSpeedEngine.__getattr__:

if name in dir(self):
    return getattr(self, name)

This eventually failed with:

RecursionError: maximum recursion depth exceeded

The updated implementation first uses the parent torch.nn.Module attribute
resolver. It delegates to the wrapped model only when the attribute is not an
engine parameter, buffer, or registered submodule.

Changes

  • Resolve registered parameters, buffers, and submodules through
    torch.nn.Module.__getattr__.
  • Preserve attribute delegation to the wrapped model.
  • Preserve normal AttributeError behavior for missing attributes.
  • Add regression coverage for:
    • registered workspace submodule lookup;
    • wrapped-model attribute delegation;
    • missing attributes.

Testing

Unit tests

pytest -q tests/unit/runtime/test_engine_attribute_delegation.py
root@6c991c12a955:/workspace/DeepSpeed_woo# pytest -q tests/unit/runtime/test_engine_attribute_delegation.py
============================================================= test session starts =============================================================
platform linux -- Python 3.12.3, pytest-9.1.1, pluggy-1.6.0 -- /usr/bin/python3.12
cachedir: .pytest_cache
rootdir: /workspace/DeepSpeed_woo/tests
configfile: pytest.ini
plugins: anyio-4.12.0
collected 3 items                                                                                                                             

tests/unit/runtime/test_engine_attribute_delegation.py::test_getattr_resolves_registered_workspace_module PASSED                        [ 33%]
tests/unit/runtime/test_engine_attribute_delegation.py::test_getattr_delegates_missing_attributes_to_model PASSED                       [ 66%]
tests/unit/runtime/test_engine_attribute_delegation.py::test_getattr_raises_for_missing_attribute PASSED                                [100%]

============================================================== warnings summary ===============================================================
unit/runtime/test_engine_attribute_delegation.py::test_getattr_resolves_registered_workspace_module
  /workspace/DeepSpeed_woo/tests/conftest.py:47: UserWarning: Running test without verifying torch version, please provide an expected torch version with --torch_ver
    warnings.warn(

unit/runtime/test_engine_attribute_delegation.py::test_getattr_resolves_registered_workspace_module
  /workspace/DeepSpeed_woo/tests/conftest.py:54: UserWarning: Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================================================== slowest durations ==============================================================

(9 durations < 1s hidden.)
======================================================= 3 passed, 2 warnings in 31.64s ========================================================
root@6c991c12a955:/workspace/DeepSpeed_woo# 

Result:

3 passed, 2 warnings in 34.15s

The warnings only reported that explicit expected Torch and CUDA versions were
not provided to the test runner.

HybridEngine integration validation

The fix was also validated through the real HybridEngine
release_inference_cache=True path using the OPSD profiling benchmark related
to #8197.

Test environment:

  • GPU: NVIDIA RTX A4500
  • Compute capability: 8.6
  • GPU memory: 20,470 MiB
  • NVIDIA driver: 580.159.04
  • PyTorch: 2.9.1+cu128
  • CUDA runtime: 12.8
  • Transformers: 4.40.2
  • Model: facebook/opt-6.7b
  • Dtype: FP16
  • Device: cuda:0
  • World size: 1
  • ZeRO stage: 0
  • Batch size: 1
  • Samples per prompt: 4
  • Prompt length: 512
  • Response length: 128
  • Warmup iterations: 5
  • Measured iterations: 20
  • CUDA_LAUNCH_BLOCKING=1

Before this fix, the first rollout failed in
retake_inference_cache() with a recursive __getattr__ traceback.

After this fix, the benchmark completed successfully and produced:

Mean total latency:      5322.63 ms
P50 total latency:       5318.74 ms
P95 total latency:       5340.41 ms
Mean generation latency: 5322.17 ms
Mean throughput:         96.19 tokens/s
Peak allocated memory:   13144.6 MB

The corresponding cache-retaining baseline was:

Mean total latency:      5189.65 ms
P50 total latency:       5185.86 ms
P95 total latency:       5215.35 ms
Mean generation latency: 5189.39 ms
Mean throughput:         98.66 tokens/s
Peak allocated memory:   13144.6 MB

For this workload, the complete cache release/retake lifecycle increased mean
rollout latency by approximately 132.98 ms, or 2.56%, and reduced aggregate
throughput by approximately 2.50%.

These measurements are included as integration validation rather than a
cross-hardware performance claim.

Scope

This PR fixes engine attribute lookup and restores the HybridEngine inference
workspace release/retake path.

It does not:

  • change the HybridEngine cache-management policy;
  • optimize cache release or reacquisition;
  • address the separate same-process, multi-shape fused-kernel error;
  • add the OPSD profiling benchmark;
  • cover ZeRO-3, LoRA fuse/unfuse, or complete OPSD step timing.

Related to #8197.

Signed-off-by: nathon-lee <leejianwoo@gmail.com>

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 49dc3e9372

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread deepspeed/runtime/engine.py
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant