[REQUEST] Add support for non-__dict__ outputs such as MinkowskiEngine SparseTensor in ZeRO Stage 3 (DeepSpeed v0.9.2)

Is your feature request related to a problem? Please describe.
When using DeepSpeed ZeRO Stage 3 (tested on DeepSpeed v0.9.2) in combination with the MinkowskiEngine, I encountered an issue during the forward pass when the model returns a MinkowskiEngine.SparseTensor as output. This object does not have a dict attribute (being a C++ extension), and DeepSpeed attempts to call vars() on it (e.g., in parameter_offload.py), leading to the following error:

TypeError: vars() argument must have dict attribute

Describe the solution you'd like
It would be great if DeepSpeed could add a safer check when handling model outputs — for example, only calling vars() on objects that have a dict, or providing a way to register custom serialization or hooks for non-standard outputs like SparseTensor.

Describe alternatives you've considered

• Avoiding ZeRO Stage 3 entirely, which sacrifices memory savings and speedups.

• Patching DeepSpeed locally, which is fragile and hard to maintain across versions.

Additional context
This issue impacts users working on sparse 3D data (e.g., scene completion) with MinkowskiEngine and DeepSpeed.
Tested on:

DeepSpeed v0.9.2 / v0.8.2

MinkowskiEngine v0.5.4

PyTorch 2.0.1 / 1.13

pytorch-lightning 2.3.0 / 1.9.0

CUDA 11.7

2 devices: RTX 3060 (12gb VRAM), RTX 3090 (24gb VRAM)

Also, I have used Stage 2, however, with the model size in the image below, the GPU RTX 3060 has been OOM, and crash the training.

Also, I cannot try mixed-precision="16" because the SparseTensor has been designed with fp32.

I would greatly appreciate it if you could suggest a possible workaround for this issue.

Issue:

Model size

Config Stage 2:

Config Stage 3:

Start by inspecting parameter_offload.py where DeepSpeed calls vars() while handling model outputs, then reproduce the ZeRO Stage 3 failure with a MinkowskiEngine SparseTensor. Determine how non-dict outputs should be handled or extended, and verify that the forward pass completes without the vars() TypeError.