fix(engine): resolve inference workspace attribute lookup

[nathon-lee]

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__:

ifnameindir(self):
returngetattr(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.

[chatgpt-codex-connector]

💡 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".

[delock]

Hi @nathon-lee thanks for the fix. I have one comments hope you can followup. Thanks!

[nathon-lee]

Hi @delock, thanks for the review. I have addressed the requested change and strengthened the regression test to verify the exact AttributeError message. The update has been pushed and is ready for another look.