369cb6c6fe
SparseGroupNorm / SparseLayerNorm used torch.zeros_like which fails with "DispatchStub: missing kernel for mps" on PyTorch builds compiled with both CUDA and MPS backends (the user's local build hit this). Added a _zeros_like_safe helper that builds zeros on CPU and transfers when the reference tensor is on MPS; Apple Silicon unified memory makes the transfer metadata-only, so the overhead vs a working MPS zeros kernel is negligible. On CPU, behaves identically to torch.zeros_like. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from ..utils import manual_cast
|
|
from . import VarLenTensor
|
|
from . import config
|
|
|
|
__all__ = [
|
|
'SparseGroupNorm',
|
|
'SparseLayerNorm',
|
|
'SparseGroupNorm32',
|
|
'SparseLayerNorm32',
|
|
]
|
|
|
|
|
|
def _zeros_like_safe(t: torch.Tensor) -> torch.Tensor:
|
|
"""zeros_like equivalent that works around PyTorch MPS builds missing
|
|
fp16/fp32 zeros kernels (DispatchStub missing). Constructs zeros on CPU
|
|
and transfers — Apple Silicon unified memory makes the transfer a
|
|
metadata-only operation, so the overhead vs a true MPS zeros kernel is
|
|
negligible when it does exist."""
|
|
if t.device.type == 'mps':
|
|
cpu_zeros = torch.zeros(t.shape, dtype=t.dtype)
|
|
return cpu_zeros.to(t.device)
|
|
return torch.zeros_like(t)
|
|
|
|
|
|
class SparseGroupNorm(nn.GroupNorm):
|
|
def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
|
|
super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine)
|
|
|
|
def forward(self, input: VarLenTensor) -> VarLenTensor:
|
|
nfeats = _zeros_like_safe(input.feats)
|
|
for k in range(input.shape[0]):
|
|
bfeats = input.feats[input.layout[k]]
|
|
bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
|
|
bfeats = super().forward(bfeats)
|
|
bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
|
|
nfeats[input.layout[k]] = bfeats
|
|
return input.replace(nfeats)
|
|
|
|
|
|
class SparseLayerNorm(nn.LayerNorm):
|
|
def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
|
|
super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine)
|
|
|
|
def forward(self, input: VarLenTensor) -> VarLenTensor:
|
|
nfeats = _zeros_like_safe(input.feats)
|
|
for k in range(input.shape[0]):
|
|
bfeats = input.feats[input.layout[k]]
|
|
bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
|
|
bfeats = super().forward(bfeats)
|
|
bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
|
|
nfeats[input.layout[k]] = bfeats
|
|
return input.replace(nfeats)
|
|
|
|
|
|
class SparseGroupNorm32(SparseGroupNorm):
|
|
"""
|
|
A GroupNorm layer that converts to float32 before the forward pass.
|
|
"""
|
|
def forward(self, x: VarLenTensor) -> VarLenTensor:
|
|
x_dtype = x.dtype
|
|
x = manual_cast(x, torch.float32)
|
|
o = super().forward(x)
|
|
return manual_cast(o, x_dtype)
|
|
|
|
|
|
class SparseLayerNorm32(SparseLayerNorm):
|
|
"""
|
|
A LayerNorm layer that converts to float32 before the forward pass.
|
|
"""
|
|
def forward(self, x: VarLenTensor) -> VarLenTensor:
|
|
x_dtype = x.dtype
|
|
x = manual_cast(x, torch.float32)
|
|
o = super().forward(x)
|
|
return manual_cast(o, x_dtype)
|