DeepSpeed

313ce47b - Fix ZeRO-3 crash in AutoTP universal-checkpoint metadata (#8270)

Fix ZeRO-3 crash in AutoTP universal-checkpoint metadata (#8270) ## Problem `DeepSpeedHybridEngine` cannot be initialized with ZeRO-3. `deepspeed.initialize()` raises before training starts: ``` File "deepspeed/runtime/hybrid_engine.py", line 354, in create_inference_module self.create_inference_containers(self.module) File "deepspeed/runtime/hybrid_engine.py", line 294, in create_inference_containers self._other_layers.append(self.inference_policies[child.__class__][0](module=child, ...)) File "deepspeed/module_inject/layers.py", line 739, in __init__ self._mark_uc_metadata() File "deepspeed/module_inject/layers.py", line 807, in _mark_uc_metadata original_weight_shape = (original_out_dim, self.weight.shape[1]) IndexError: tuple index out of range ``` ZeRO-3 + HybridEngine is the standard DeepSpeed-Chat actor configuration, so this blocks the engine's primary use case. ## Cause `_mark_uc_metadata()` was added in #7908 (universal checkpoint for AutoTP) and runs unconditionally from `LinearLayer.__init__` / `LinearAllreduce.__init__`. It reads `param.shape[1]` to reconstruct the pre-TP parameter shape. Under ZeRO-3 a partitioned parameter's local data is an empty 1-D tensor, so that index is out of range. Observed on `facebook/opt-1.3b`, ZeRO-3, single GPU: ``` LinearLayer weight: shape=(0,) numel=0 ds_id=0 ds_shape=torch.Size([50272, 2048]) ds_status=ZeroParamStatus.NOT_AVAILABLE tp_world=1 ``` HybridEngine constructs a `LinearLayer` for every non-transformer layer (embeddings, `lm_head`, final norm) regardless of TP size, so it hits this path even with `inference_tp_size=1`. ## Fix Read the pre-partition shape from `ds_shape` instead of the local `.shape`. ZeRO-3 sets `param.ds_shape = param.shape` in `partition_parameters.py` *before* flattening the parameter's data, so `ds_shape` is exactly the value these call sites were already trying to read. When ZeRO-3 is not in use the attribute is absent and `param.shape` is used, so the existing AutoTP path is unchanged. The helper lives on `TensorParallel_Layer` and is used by both the column-parallel (`LinearLayer`) and row-parallel (`LinearAllreduce`) implementations, which had the same bug. `SubParamLinearLayer` / `SubParamLinearAllreduce` take their shapes from precomputed `_logical_shape` / `_orig_weight_shape` attributes rather than indexing `param.shape`, so they are not affected and are left alone. ## Tests Two regression tests added to `tests/unit/runtime/tensor_parallel/test_autotp_universal_checkpoint.py`, covering both the column-parallel and row-parallel paths. They construct a parameter in the state ZeRO-3 leaves it in and assert the recorded metadata carries the pre-partition shape. Both fail on master with the exact `IndexError` (at lines 707 and 807) and pass with this change. The 10 existing tests in that file continue to pass. ## Verification End-to-end RLHF-style loop (`eval()` → `generate()` → `train()` → `forward`/`backward`/`step()`), `facebook/opt-1.3b`, batch 8, 256-token prompt, 128 new tokens, H100: - **ZeRO-3** — fails on master at `deepspeed.initialize()`; runs to completion with this change (2.63 s/iter steady state). - **ZeRO-2** — unaffected (796 ms/iter before and after). Signed-off-by: Zhipeng Wang <zhipeng.rainbowserie@gmail.com>