fix(inference): Add missing dtype attribute to ParameterBase setter (#7378)

### Description

This PR fixes an `AttributeError: 'UnembedParameter' object has no
attribute 'dtype'` that occurs in the Inference V2 engine. The issue is
triggered when using a high-level interface like
[DeepSpeed-MII](https://github.com/deepspeedai/DeepSpeed-MII) to run
inference on models with tied input/output embeddings, such as Llama 2.

**Resolves: #7260**

### Root Cause Analysis

The root cause is that while the `ParameterBase` metaclass correctly
creates property setters for parameter tensors, the setter function
(`param_setter`) only assigns the tensor value itself. It does not
propagate the tensor's `dtype` to the container instance.

Downstream functions, such as `flatten_inference_model`, expect every
parameter container to have a `.dtype` attribute. When they encounter a
custom container like `UnembedParameter` that lacks this attribute, an
`AttributeError` is raised.

### The Fix

The solution is to modify the `param_setter` function within
`make_param_setter` located in
`deepspeed/inference/v2/model_implementations/parameter_base.py`.

I have added the line `self.dtype = value.dtype` immediately after the
parameter tensor is assigned. This simple change ensures that any object
inheriting from `ParameterBase` will now correctly expose the `dtype` of
the tensor it wraps, resolving the error.

### Verification

This fix has been thoroughly verified in a containerized GPU environment
(RunPod with PyTorch 2.1). The verification process involved:
1. Cloning both the `deepspeed` and `DeepSpeed-MII` repositories from
source.
2. Installing the modified `deepspeed` library from this branch.
3. Installing the `DeepSpeed-MII` library (with a packaging fix) to
trigger the bug.
4. Running an end-to-end inference script with `mii.pipeline` and a
standard language model.

The logs confirm that with this fix, the program successfully executes
past the original point of failure. The `AttributeError` is completely
resolved, and the DeepSpeed engine proceeds correctly to the model
loading phase.

*(Note: A full end-to-end run in the test environment was ultimately
blocked by a separate, pre-existing build issue in DeepSpeed's op
builder (`ModuleNotFoundError: dskernels`), which is unrelated to this
logic fix. The successful progression past the original error point
serves as definitive proof of this fix's effectiveness.)*

### Related Context

This bug is primarily triggered via the
[**DeepSpeed-MII**](https://github.com/deepspeedai/DeepSpeed-MII)
project. A companion PR,
**[deepspeedai/DeepSpeed-MII#567](https://github.com/deepspeedai/DeepSpeed-MII/pull/567)**,
has been submitted to fix a packaging issue in that repository that was
a prerequisite for this verification.

output:

<img width="1014" alt="Screenshot 2025-06-22 at 14 16 15"
src="https://github.com/user-attachments/assets/1a658f98-a98b-4584-ae11-59e9edfd0b7e"
/>

<img width="1012" alt="Screenshot 2025-06-22 at 14 16 26"
src="https://github.com/user-attachments/assets/3959d0e5-d6dc-4ed4-adbc-6919e00da172"
/>

<img width="1728" alt="Screenshot 2025-06-22 at 14 17 40"
src="https://github.com/user-attachments/assets/537fd354-b840-4af2-98ab-d243c6902412"
/>

Signed-off-by: Vensenmu <vensenmu@gmail.com>
Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
This commit is contained in:
Vensen
2025-06-24 07:38:35 +08:00
committed by GitHub
parent 2a450b3a33
commit 61829b55ea

View File

@ -35,6 +35,7 @@ def make_param_setter(clsname, param):
def param_setter(self, value):
setattr(self, f"__{clsname}__{param}", value)
self.dtype = value.dtype
self.complete_component()
return param_setter