Fix MLU_Accelerator conformance to the DeepSpeedAccelerator ABC (#8208)
`MLU_Accelerator` diverges from the `DeepSpeedAccelerator` contract in
three places. I compared all nine backends against the ABC and it is the
only one that does.
## 1. `pin_memory` drops `align_bytes`
(`accelerator/mlu_accelerator.py:226`)
The ABC declares `pin_memory(self, tensor, align_bytes=1)`
(`abstract_accelerator.py:264`); MLU declares `pin_memory(self,
tensor)`. Three in-tree call sites pass it by keyword, all on the
ZeRO-Infinity / NVMe offload swap path:
- `deepspeed/runtime/swap_tensor/utils.py:188`
- `deepspeed/runtime/swap_tensor/partitioned_param_swapper.py:131`
- `deepspeed/runtime/swap_tensor/partitioned_param_swapper.py:398`
all `align_bytes=0`, so on MLU these raise `TypeError` rather than
differing quietly.
## 2. `initial_seed` requires an argument the ABC does not declare
(`:81`)
#5569 changed `initial_seed(self, seed)` to `initial_seed(self)` in the
ABC and in cpu/cuda/hpu/mps/npu/xpu, on the grounds that
`torch.<device>.initial_seed()` takes no argument. MLU landed 3 months
later in #6472 as a copy of the pre-#5569 NPU code, which brought the
argument back, so `get_accelerator().initial_seed()` raises `TypeError`
on MLU alone.
To be accurate about severity: the only in-tree reference is commented
out (`deepspeed/runtime/pipe/module.py:200`), so this is a break in the
public accelerator API and a regression of the #5569 decision, not a
live in-tree crash.
## 3. `create_graph` builds a graph and drops it (`:186`)
```python
def create_graph(self):
torch.mlu.MLUGraph()
```
The in-tree consumers use the returned handle: `graph_process()`
(`deepspeed/runtime/utils.py:101`) passes it straight into
`capture_to_graph()` and then `replay_graph()`, so both get `None`.
Worth separating from a deliberate opt-out. cpu, mps, npu, sdaa and xpu
return `None` *and* pair it with a `noop_context()` `capture_to_graph`.
MLU pairs the missing return with a real capture
(`torch.mlu.graph(...)`, `:189`) and a real `graph.replay()` (`:192`),
which is the cuda/hpu/supa shape. That reads as an oversight rather than
an abstention, but you would know better than I would.
## The fix
Three one-line changes, matching what the other backends already do: add
`align_bytes=1` to `pin_memory` (the cuda/npu/sdaa/supa body is
byte-identical to MLU's, so nothing else changes), drop the `seed`
argument from `initial_seed`, and return from `create_graph`.
## How I verified
- Added two CPU-only tests to
`tests/unit/accelerator/test_accelerator.py`: a Liskov check that an
override may widen the abstract signature but never narrow it, and a
check that `create_graph` returns. On `df84f6d8` the file goes from 2
failed / 25 passed to 27 passed; both failures are MLU and the other
eight backends pass untouched.
- The signature check is deliberately not "arity must match": every
backend gives `device`/`device_name` a default the ABC declares
required, and that widening is fine. Only requiring *more* than the ABC,
or dropping one of its defaulted parameters, fails. That is what keeps
the cpu/mps `create_op_builder` naming difference from tripping it.
- Both tests need no accelerator hardware, so they run in the existing
`cpu-torch-latest` job. yapf 0.40.0 reports both files clean.
- What I did not check: I have no Cambricon hardware, so none of this
ran on an MLU device. The three contract faults are exercised directly,
but the downstream consequences I describe (the swap path,
`graph_process`) are read from the call sites rather than observed on
silicon.
Happy to split this into three commits, or to drop the tests if you
would rather keep the suite as it is.
Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>