Repository: pytorch
Issue: [Compiled_autograd] running nn.LayerNorm failed for torch.compile with compiled_autograd when deepspeed Zero3
Labels: oncall: distributed,triaged,oncall: pt2,module: compiled autograd

### 🐛 Describe the bug

When running a simple model including torch.nn.LayerNorm using deepspeed zero3 with torch.compile and [compiled_autograd](https://github.com/pytorch/tutorials/blob/main/intermediate_source/compiled_autograd_tutorial.rst). An error occurs:

> site-packages/torch/_subclasses/fake_tensor.py:2017] RuntimeError: Attempting to broadcast a dimension of length 0 at -1! Mismatching argument at index 1 had torch.Size([0]); but expected shape should be broadcastable to [100, 120]

We first found this error in BERT model with deepspeed Zero3 with torch.compile and compiled_autograd.

- It's ok for deepspeed Zero1/2 with torch.compile and compiled_autograd
- It's ok for deepspeed Zero3 with torch.compile and without compiled_autograd
- There are a lot of graph beaks and recompiles in deepspeed Zero3 with torch.compile. And Zero3 will partition model parameters through [hooks](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/zero/parameter_offload.py#L259). 
- To simplify the issue, I made a small reproducer to extract error op(torch.nn.LayerNorm)

**Investigation**
The error: "RuntimeError: Attempting to broadcast a dimension of length 0 at -1! Mismatching argument at index 1 had torch.Size([0]); but expected shape should be broadcastable to [100, 120]"
**It occurs when compiled autograd tries to trace the backward graph.**
It appears in [LayerNorm backward decompositions](https://github.com/pytorch/pytorch/blob/main/torch/_decomp/decompositions.py#L1703). It tries to broadcast weight_cast(torch.Size([0]) to grad_out_cast' shape([100, 120]) and fails.
```
if weight_cast is not None:         
    grad_x_hat = grad_out_cast * weight_cast 
```
If bypassing the LayerNorm weight by setting `nn.LayerNorm(120, eps=1e-12, elementwise_affine=False)` instead of `elementwise_affine=True` in the file deepspeed_reproducer_cpu.py, the running is ok.

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @chauhang @penguinwu @xmfan @yf225 @ezyang @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 

### Error logs

```
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home1/yitingw1/habana/deepspeed_demo/deepspeed_reproducer_cpu.py", line 83, in <module>
[rank0]:     model_engine.backward(loss)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
[rank0]:     ret_val = func(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2020, in backward
[rank0]:     self.optimizer.backward(loss, retain_graph=retain_graph)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn
[rank0]:     ret_val = func(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/deepspeed/runtime/zero/stage3.py", line 2250, in backward
[rank0]:     self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
[rank0]:     scaled_loss.backward(retain_graph=retain_graph)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_tensor.py", line 581, in backward
[rank0]:     torch.autograd.backward(
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/autograd/__init__.py", line 347, in backward
[rank0]:     _engine_run_backward(
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/utils/_stats.py", line 21, in wrapper
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 1308, in __torch_dispatch__
[rank0]:     return proxy_call(self, func, self.pre_dispatch, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 906, in proxy_call
[rank0]:     out = func(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_ops.py", line 716, in __call__
[rank0]:     return self._op(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/utils/_stats.py", line 21, in wrapper
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1238, in __torch_dispatch__
[rank0]:     return self.dispatch(func, types, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1692, in dispatch
[rank0]:     return self._cached_dispatch_impl(func, types, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1348, in _cached_dispatch_impl
[rank0]:     output = self._dispatch_impl(func, types, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1943, in _dispatch_impl
[rank0]:     return decomposition_table[func](*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_decomp/decompositions.py", line 1729, in native_layer_norm_backward
[rank0]:     grad_x_hat = grad_out_cast * weight_cast
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/utils/_stats.py", line 21, in wrapper
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1238, in __torch_dispatch__
[rank0]:     return self.dispatch(func, types, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1692, in dispatch
[rank0]:     return self._cached_dispatch_impl(func, types, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 1339, in _cached_dispatch_impl
[rank0]:     output = self._dispatch_impl(func, types, args, kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_subclasses/fake_tensor.py", line 2013, in _dispatch_impl
[rank0]:     r = func(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_ops.py", line 716, in __call__
[rank0]:     return self._op(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 273, in _fn
[rank0]:     result = fn(*args, **kwargs)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_prims_common/wrappers.py", line 141, in _fn
[rank0]:     result = fn(**bound.arguments)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_refs/__init__.py", line 1049, in _ref
[rank0]:     a, b = _maybe_broadcast(a, b)
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_refs/__init__.py", line 422, in _maybe_broadcast
[rank0]:     common_shape = _broadcast_shapes(
[rank0]:   File "/home1/yitingw1/mambaforge/envs/wyt_pt/lib/python3.10/site-packages/torch/_refs/__init__.py", line 411, in _broadcast_shapes
[rank0]:     raise RuntimeError(
[rank0]: RuntimeError: Attempting to broadcast a dimension of length 0 at -1! Mismatching argument at index 1 had torch.Size([0]); but expected shape should be broadcastable to [100, 120]
```

### Minified repro

Running script:
`TORCHDYNAMO_EXTENDED_DEBUG_CPP=1 TORCH_LOGS="+dynamo,graph,graph_code,graph_breaks,recompiles,aot_graphs,aot_joint_graph,compiled_autograd_verbose"
deepspeed --num_nodes 1 --num_gpus 1 deepspeed_reproducer_cpu.py `

Below is deepspeed_reproducer_cpu.py
```python
import torch
import torchvision
import torchvision.transforms as transforms
import torch.distributed as dist
import deepspeed
from deepspeed.accelerator import get_accelerator
from tqdm import tqdm
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(32 * 32 * 3, 120)
        self.fc2 = nn.Linear(120, 10)
        self.LayerNorm1 = nn.LayerNorm(120, eps=1e-12, elementwise_affine=True)

    def forward(self, x):
        x = torch.flatten(x, 1)  # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = self.LayerNorm1(x)
        x = self.fc2(x)
        return x

compile_kwargs = {"dynamic": False}
device = torch.device('cpu')

model = Net()
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
model_engine, optimizer, *_ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    optimizer=optimizer,
    config="./deepspeed_config.json",
)
# torch_compile
model_engine.compile(
    compile_kwargs=compile_kwargs,
)

# dataset
transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
)
batch_size = 100
trainset = torchvision.datasets.CIFAR10(
    root="./DATA/CIFAR10", train=True, download=True, transform=transform
)
# process dataset
trainloader = DataLoader(
    trainset,
    batch_size=batch_size,
    sampler=DistributedSampler(trainset, shuffle=True),
    num_workers=16,
    pin_memory=True,
)
progress_bar = tqdm(
    total=len(trainloader),
    desc=f"Training 1/1 epoch",
    position=0,
    leave=True,
    disable= dist.is_initialized() and dist.get_rank() != 0,
)
for epoch in range(100):
    with torch._dynamo.compiled_autograd.enable(
                torch.compile(backend=get_accelerator().get_compile_backend(), **compile_kwargs)):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            
            # forward + backward + optimize
            outputs = model_engine(inputs)
            loss = criterion(outputs, labels)
            model_engine.backward(loss)
            model_engine.step()

            # print statistics
            running_loss += loss.item()
            if i % 2000 == 1999:  # print every 2000 mini-batches
                print(f"[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}")
                running_loss = 0.0
            progress_bar.update(1)
print("Finished Training")
```

Below is deepspeed_config.json
```json
{
    "train_batch_size": 32, 
    "optimizer": {
        "type": "SGD",
        "params": {
            "lr": 0.001,
            "momentum": 0.9
        }
    },
    "zero_allow_untested_optimizer": true,
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": false,
        "reduce_scatter" : false,
        "contiguous_gradients" : false
    },
}
```

### Versions

Collecting environment information...
PyTorch version: 2.5.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.27.4
Libc version: glibc-2.35

Python version: 3.10.0 | packaged by conda-forge | (default, Nov 20 2021, 02:24:10) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-102-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.26.3
[pip3] torch==2.5.1+cpu
[pip3] torchaudio==2.5.1+cpu
[pip3] torchvision==0.20.1+cpu
[pip3] deepspeed==0.15.3