DeepSpeed

1707800a - Assert offload-state removal only when offload leaves the accelerator (#8738)

Assert offload-state removal only when offload leaves the accelerator (#8738) ## Description `validate_device` in `tests/unit/v1/zero/test_offload_states.py` requires the contiguous grad buffer to be **gone** once it is reported on cpu, because a cpu device means `offload_states()` moved it to the host tier: ```python if state == OffloadStateTypeEnum.contiguous_grad_buffer and state_device[state] == torch.device("cpu"): assert len(get_state_devices(model, state)) == 0, f"State {state} must be removed after offload_states()" ``` On the cpu accelerator the offload target **is** the accelerator, so `state_device[state] == torch.device("cpu")` holds for everything, yet the buffer was never moved anywhere — the assertion fails for every parametrization that includes it. This is the device-state counterpart of #8409, which gated the **memory-delta** assertions on `allocator_backed_memory_stats` for the same reason. Add a sibling `offload_leaves_accelerator` flag (next to that one) and require it for the removal check. ## Behavior | Accelerator | `offload_leaves_accelerator` | Effect | |---|---|---| | CUDA / XPU | `True` | assertion unchanged, bit for bit | | cpu | `False` | state falls through to the device comparison instead, which holds | ## Validation (executed on real hardware) 20-core x86_64 CPU without AVX512-FP16, gloo backend, 2 ranks (`LOCAL_SIZE=2`), `-n 4`: - **Without** this change on master: `16 failed, 112 passed` — and all 16 failures are this single assertion (`State OffloadStateTypeEnum.contiguous_grad_buffer must be removed after offload_states()`). - **With** this change: `128 passed, 0 failed`. So the cpu accelerator keeps the coverage of the other 112 parametrizations instead of losing them to a blanket skip. The enclosing cpu CI run is #8381, which exercises this file at multi-rank and reports the same result (`128 PASSED`). Signed-off-by: Guokai Ma <guokai.ma@intel.com>