[BUG] Deepspeed 0.19.8 fails with 0-dimensional nn.Parameter

Il traceback indica DeepSpeedEngine._broadcast_model in deepspeed/runtime/engine.py, dove p.data.view(torch.uint8) non riesce con un parametro scalare. Inizia leggendo il codice di broadcast e la modifica in #8416, poi verifica se ci sono test che coprono i broadcast dei modelli con parametri 0-D. Il lavoro è completato quando la riproduzione fornita si inizializza correttamente con il parametro scalare.

Describe the bug
Deepspeed 0.19.8 fails with 0-dimensional nn.Parameter.
This bug was introduced in #8416 due to param.data.view(torch.uint8), as PyTorch doesn't allow viewing scalar tensor as another data type with different element size

To Reproduce

import contextlib
import os
import torch
import torch.distributed as dist
import deepspeed as ds


@contextlib.contextmanager
def init_process_group():
    LOCAL_RANK = os.getenv("LOCAL_RANK", None)

    try:
        # Create process group of size 2
        dist.init_process_group(
            backend="nccl", store=dist.HashStore(), world_size=1, rank=0
        )
        os.environ["LOCAL_RANK"] = "0"
        yield dist.new_group(ranks=[0])
    finally:
        if dist.is_initialized():
            dist.destroy_process_group()
        if LOCAL_RANK is None:
            del os.environ["LOCAL_RANK"]
        else:
            os.environ["LOCAL_RANK"] = LOCAL_RANK


def deepspeed_zero3_offload_config():
    return {
        "zero_optimization": {
            "stage": 3,
            "offload_param": {"device": "cpu"},
            "offload_optimizer": {"device": "cpu"},
        },
        "train_batch_size": 1,
        "optimizer": {
            "type": "AdamW",
            "params": {
                "lr": 1e-2,
                "betas": [0.9, 0.999],
                "eps": 1e-7,
                "weight_decay": 1e-5,
            },
        },
        "gradient_clipping": 0.0,
    }


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 10)
        self.scalar = torch.nn.Parameter(torch.tensor(1.0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(x) * self.scalar


def main():
    model = Model().cuda().eval()

    with init_process_group():
        # ds.initialize fails with the following error:
        #   [rank0]: Traceback (most recent call last):
        #   [rank0]:   File "main.py", line 70, in <module>
        #   [rank0]:     main()
        #   [rank0]:   File "main.py", line 61, in main
        #   [rank0]:     _ = ds.initialize(
        #   [rank0]:   File "<...>/python3.10/site-packages/deepspeed/__init__.py", line 569, in initialize
        #   [rank0]:     engine = DeepSpeedEngine(args=args,
        #   [rank0]:   File "<...>/python3.10/site-packages/deepspeed/runtime/engine.py", line 621, in __init__
        #   [rank0]:     self._configure_distributed_model(model)
        #   [rank0]:   File "<...>/python3.10/site-packages/deepspeed/runtime/engine.py", line 2195, in _configure_distributed_model
        #   [rank0]:     self._broadcast_model()
        #   [rank0]:   File "<...>/python3.10/site-packages/deepspeed/runtime/engine.py", line 2059, in _broadcast_model
        #   [rank0]:     dist.broadcast(p.data.view(torch.uint8),
        #   [rank0]: RuntimeError: self.dim() cannot be 0 to view Float as Byte (different element sizes)
        _ = ds.initialize(
            model=model,
            model_parameters=list(model.parameters()),
            config=deepspeed_zero3_offload_config(),
        )


if __name__ == "__main__":
    main()

Expected behavior
Should handle 0-D scalar parameters normally, as in deepspeed 0.19.7 or earlier

ds_report output

--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
pin_memory ............. [NO] ....... [OKAY]
dc ..................... [NO] ....... [OKAY]
 [WARNING]  Unable to locate CUTLASS. Install nvidia-cutlass, clone CUTLASS next to DeepSpeed, or set CUTLASS_PATH to the CUTLASS checkout.
evoformer_attn ......... [NO] ....... [NO]
fp_quantizer ........... [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
 [WARNING]  gds requires the dev libaio .so object and headers but these were not found.
 [WARNING]  gds: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
gds .................... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
spatial_inference ...... [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/local/mnt/workspace/kyunggeu/miniconda3/envs/aimet-torch-dev/lib/python3.10/site-packages/torch']
torch version .................... 2.14.0+cu126
deepspeed install path ........... ['/local/mnt/workspace/kyunggeu/miniconda3/envs/aimet-torch-dev/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.19.8, unknown, unknown
torch cuda version ............... 12.6
torch hip version ................ None
nvcc version ..................... 12.2
deepspeed wheel compiled w. ...... torch 0.0, cuda 0.0
shared memory (/dev/shm) size .... 15.51 GB