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

[Flink-ddd]

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 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 project. A companion PR, deepspeedai/DeepSpeed-MII#567, has been submitted to fix a packaging issue in that repository that was a prerequisite for this verification.

output：

[tohtana]

Thank you @Flink-ddd!

[Flink-ddd]

Hi @tohtana ,
Thank you so much for your time in reviewing and merging this PR! I truly appreciate it.

I just wanted to provide a little more context on the companion PR for the DeepSpeed-MII project, as it was a necessary prerequisite for verifying this fix. The PR is: deepspeedai/DeepSpeed-MII#567.

The change to setup.py in that PR was essential. During testing in a clean containerized environment, the original find_packages() call in DeepSpeed-MII's setup script unreliably failed to discover the mii package, leading to a persistent ModuleNotFoundError. Making the package discovery explicit with include=['mii', 'mii.*'] resolves this build issue, ensuring a robust installation from source and improving the developer experience.

If you have a moment, I would be grateful if you could take a look. No rush at all, and thanks again for maintaining this great project!