DeepSpeed
437226bf - Fix MLU_Accelerator conformance to the DeepSpeedAccelerator ABC (#8208)

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