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

Commit
8 days ago
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>
Author
Parents
Loading