fix(engine): resolve inference workspace attribute lookup (#8288)
## 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:
```python
"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__`:
```python
if name in dir(self):
return getattr(self, name)
```
This eventually failed with:
```text
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
```bash
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:
```text
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:
```text
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:
```text
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>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>