Skip to content

Patch training bfloat16 on MPS bug #1278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mrariden merged 2 commits into from
Jul 23, 2025
+13 −1
Merged

Patch training bfloat16 on MPS bug #1278

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7e545ec
warn when training with bfloat16 on MPS; fallback to float32
mrariden Jul 21, 2025
f2b98f4
revert back to originial dtype after training
mrariden Jul 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 13 additions & 1 deletion cellpose/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -319,7 +319,7 @@ def train_seg(net, train_data=None, train_labels=None, train_files=None,
Train the network with images for segmentation.

Args:
net (object): The network model to train.
net (object): The network model to train. If `net` is a bfloat16 model on MPS, it will be converted to float32 for training. The saved models will be in float32, but the original model will be returned in bfloat16 for consistency. CUDA/CPU will train in bfloat16 if that is the provided net dtype.
train_data (List[np.ndarray], optional): List of arrays (2D or 3D) - images for training. Defaults to None.
train_labels (List[np.ndarray], optional): List of arrays (2D or 3D) - labels for train_data, where 0=no masks; 1,2,...=mask labels. Defaults to None.
train_files (List[str], optional): List of strings - file names for images in train_data (to save flows for future runs). Defaults to None.
Expand Down Expand Up @@ -357,6 +357,14 @@ def train_seg(net, train_data=None, train_labels=None, train_files=None,

device = net.device

original_net_dtype = None
if device.type == 'mps' and net.dtype == torch.bfloat16:
# NOTE: this produces a side effect of returning a network that is not of a guaranteed dtype \
original_net_dtype = torch.bfloat16
train_logger.warning("Training with bfloat16 on MPS is not supported, using float32 network instead")
net.dtype = torch.float32
net.to(torch.float32)

scale_range = 0.5 if scale_range is None else scale_range

if isinstance(normalize, dict):
Expand Down Expand Up @@ -533,4 +541,8 @@ def train_seg(net, train_data=None, train_labels=None, train_files=None,

net.save_model(filename)

if original_net_dtype is not None:
net.dtype = original_net_dtype
net.to(original_net_dtype)

return filename, train_losses, test_losses
Loading