From 3173659c2d6347585bff69cf70dbdbe02a4387a6 Mon Sep 17 00:00:00 2001 From: fnachon Date: Sat, 10 Jan 2026 15:54:29 +0100 Subject: [PATCH 1/6] Minor adaptations for Mac GPU MPS compatibility Changes made to run without errors on the Mac MPS device: torch.autocast, number of devices and workers to use on M1-5 chips, workaround for CUDA-specific code, handling of float64 incompatibilities for MPS. --- README.md | 16 +++++++++++-- pyproject.toml | 12 +++++----- src/boltzgen/cli/boltzgen.py | 8 +++---- src/boltzgen/model/loss/diffusion.py | 14 ++++++++--- src/boltzgen/model/models/boltz.py | 23 +++++++++++++------ src/boltzgen/model/modules/trunk.py | 4 ++-- src/boltzgen/model/validation/refolding.py | 9 +++++--- .../task/predict/data_from_generated.py | 2 ++ src/boltzgen/task/predict/data_from_yaml.py | 3 +++ src/boltzgen/task/predict/data_ligands.py | 2 ++ .../task/predict/data_protein_binder.py | 2 ++ 11 files changed, 68 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index fdff22c2..1b1ea1de 100755 --- a/README.md +++ b/README.md @@ -24,7 +24,7 @@ pip install boltzgen Choose the installer for your operating system, download it, and follow the on-screen prompts: * **Windows:** -* **macOS / Linux:** +* **MacOS / Linux:** After installation, **open a terminal / command prompt** (you may need to search for “Anaconda Prompt” on Windows). @@ -35,6 +35,17 @@ Run the command below in a terminal to create a fresh environment called `bg` wi ```bash conda create -n bg python=3.12 ``` +* **MacOS** + +Create a new conda environment for boltzgen with python 3.12, numba, numpy and lvmlite: + +``` +conda create --name bg python=3.12 llvmlite==0.44.0 numba==0.61.0 numpy==2.0.2 +``` +Temporary fix for loading multiple libomp +``` +export KMP_DUPLICATE_LIB_OK=TRUE +``` ### 3 - Activate the environment (do this every time you work with BoltzGen) @@ -94,7 +105,8 @@ docker build -t boltzgen:weights --build-arg DOWNLOAD_WEIGHTS=true . `boltzgen run` takes a [design specification](#how-to-make-a-design-specification-yaml) `.yaml` and produces a set of ranked designs.\ ⚠️ it downloads models (~6GB) to `~/.cache`. This can by changed by passing `--cache YOUR_PATH` or by setting `$HF_HOME`.\ -⚠️ If your run is ever interrupted, you can restart it with `--reuse`. No progress is lost. +⚠️ If your run is ever interrupted, you can restart it with `--reuse`. No progress is lost.\ +⚠️ On MacOS set `--num_workers 0` to prevent MPS-related incompatibilities and runtime errors. Optionnaly, to suppress MPS warnings about unsupported pinned memory, set `--config data.pin_memory=false` for the following steps: `design`, `inverse_folding`, `folding`, and `design_folding` ```bash diff --git a/pyproject.toml b/pyproject.toml index 21e769ab..1be8dba2 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,8 +14,8 @@ readme = { file = "PYPI_DESCRIPTION.md", content-type = "text/markdown" } description = "Protein design" dependencies = [ # Add runtime dependencies here - "numpy==2.0.2", - "numba==0.61.0", + "numpy==2.0.2; platform_system != 'Darwin'", + "numba==0.61.0; platform_system != 'Darwin'", "matplotlib", "hydride", "biotite", @@ -33,10 +33,10 @@ dependencies = [ "einx", "einops", "mashumaro", - "nvidia-ml-py>=12.535.133", - "cuequivariance_ops_cu12>=0.5.0", - "cuequivariance_ops_torch_cu12>=0.5.0", - "cuequivariance_torch>=0.5.0", + "nvidia-ml-py>=12.535.133; platform_system != 'Darwin'", + "cuequivariance_ops_cu12>=0.5.0; platform_system != 'Darwin'", + "cuequivariance_ops_torch_cu12>=0.5.0; platform_system != 'Darwin'", + "cuequivariance_torch>=0.5.0; platform_system != 'Darwin'", "huggingface_hub", "biopython", ] diff --git a/src/boltzgen/cli/boltzgen.py b/src/boltzgen/cli/boltzgen.py index 1a7f8f17..4a481431 100644 --- a/src/boltzgen/cli/boltzgen.py +++ b/src/boltzgen/cli/boltzgen.py @@ -907,8 +907,8 @@ def __init__(self, args: argparse.Namespace, moldir: Path): f"Invalid protocol: {protocol}. Valid protocols: {list(protocol_configs.keys())}" ) - # Handle use_kernels argument - device_capability = torch.cuda.get_device_capability() + # Handle use_kernels argument, defaulting to (0,0) for CPU/MPS + device_capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else (0, 0) use_kernels = None if args.use_kernels == "auto": use_kernels = device_capability[0] >= 8 @@ -927,9 +927,9 @@ def __init__(self, args: argparse.Namespace, moldir: Path): config_args_by_step = parse_config_args( protocol_config, args.config, step_names ) - + # Determine number of devices to use, defaulting to 1 for MPS/CPU devices = ( - args.devices if args.devices is not None else torch.cuda.device_count() + args.devices if args.devices is not None else torch.cuda.device_count() if torch.cuda.is_available() else 1 ) print(f"Using {devices} devices") diff --git a/src/boltzgen/model/loss/diffusion.py b/src/boltzgen/model/loss/diffusion.py index b39705ad..f8b091ba 100755 --- a/src/boltzgen/model/loss/diffusion.py +++ b/src/boltzgen/model/loss/diffusion.py @@ -80,9 +80,17 @@ def weighted_rigid_align( original_dtype = cov_matrix.dtype cov_matrix_32 = cov_matrix.to(dtype=torch.float32) - U, S, V = torch.linalg.svd( - cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None - ) + # move cov_matrix_32 to cpu for mps compatibility + if cov_matrix_32.device.type == "mps": + cov_matrix_cpu = cov_matrix_32.cpu() + U, S, V = torch.linalg.svd(cov_matrix_cpu, driver=None) + U = U.to(cov_matrix_32.device) + S = S.to(cov_matrix_32.device) + V = V.to(cov_matrix_32.device) + else: + U, S, V = torch.linalg.svd( + cov_matrix_32, driver="gesvd" if cov_matrix_32.is_cuda else None + ) V = V.mH # Catch ambiguous rotation by checking the magnitude of singular values diff --git a/src/boltzgen/model/models/boltz.py b/src/boltzgen/model/models/boltz.py index 8214b9eb..0ea75100 100755 --- a/src/boltzgen/model/models/boltz.py +++ b/src/boltzgen/model/models/boltz.py @@ -655,7 +655,7 @@ def forward( ): if self.inference_logging: print("\nRunning Structure Module.\n") - with torch.autocast("cuda", enabled=False): + with torch.autocast(device_type=s.device.type, enabled=False): if not self.inverse_fold: struct_out = self.structure_module.sample( s_trunk=s.float(), @@ -711,7 +711,7 @@ def forward( feats["coords"] = atom_coords # (multiplicity, L, 3) assert len(feats["coords"].shape) == 3 - with torch.autocast("cuda", enabled=False): + with torch.autocast(device_type=atom_coords.device.type, enabled=False): if not self.inverse_fold: struct_out = self.structure_module( s_trunk=s.float(), @@ -769,7 +769,7 @@ def forward( ] s_inputs = self.input_embedder(feats, affinity=True) - with torch.autocast("cuda", enabled=False): + with torch.autocast(device_type=s_inputs.device.type, enabled=False): if self.affinity_ensemble: dict_out_affinity1 = self.affinity_module1( s_inputs=s_inputs.detach(), @@ -1112,7 +1112,7 @@ def parameter_norm(self, module): parameters = [p.norm(p=2) ** 2 for p in module.parameters() if p.requires_grad] if len(parameters) == 0: return torch.tensor( - 0.0, device="cuda" if torch.cuda.is_available() else "cpu" + 0.0, device="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" ) norm = torch.stack(parameters).sum().sqrt() return norm @@ -1165,7 +1165,10 @@ def validation_step( "res_type =", batch["res_type"].shape, ) - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() return raise e else: @@ -1184,7 +1187,10 @@ def validation_step( if "out of memory" in str(e): msg = f"| WARNING: ran out of memory, skipping batch, {idx_dataset}" print(msg) - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() return raise e @@ -1368,7 +1374,10 @@ def predict_step( except RuntimeError as e: # catch out of memory exceptions if "out of memory" in str(e): print("| WARNING: ran out of memory, skipping batch") - torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() return {"exception": True} else: raise e diff --git a/src/boltzgen/model/modules/trunk.py b/src/boltzgen/model/modules/trunk.py index a00b0400..0858a932 100755 --- a/src/boltzgen/model/modules/trunk.py +++ b/src/boltzgen/model/modules/trunk.py @@ -355,7 +355,7 @@ def forward( asym_mask = asym_mask[:, None].expand(-1, T, -1, -1) # Compute template features - with torch.autocast(device_type="cuda", enabled=False): + with torch.autocast(device_type=cb_coords.device.type, enabled=False): # Compute distogram cb_dists = torch.cdist(cb_coords, cb_coords) boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1) @@ -504,7 +504,7 @@ def forward( token_coords = feats["center_coords"] # Compute template features - with torch.autocast(device_type="cuda", enabled=False): + with torch.autocast(device_type=token_coords.device.type, enabled=False): # Compute distogram dists = torch.cdist(token_coords, token_coords) boundaries = torch.linspace(self.min_dist, self.max_dist, self.num_bins - 1) diff --git a/src/boltzgen/model/validation/refolding.py b/src/boltzgen/model/validation/refolding.py index abf3a5ce..ea9faed1 100755 --- a/src/boltzgen/model/validation/refolding.py +++ b/src/boltzgen/model/validation/refolding.py @@ -299,11 +299,14 @@ def on_epoch_end(self, model): self.folding_model = None del self.affinity_model self.affinity_model = None - torch._C._cuda_clearCublasWorkspaces() + if torch.cuda.is_available(): + torch._C._cuda_clearCublasWorkspaces() torch._dynamo.reset() gc.collect() - torch.cuda.empty_cache() - + if torch.cuda.is_available(): + torch.cuda.empty_cache() + elif torch.backends.mps.is_available(): + torch.mps.empty_cache() # Compute standard metrics self.common_on_epoch_end(model, logname="val_monomer_ligand") self.on_epoch_end_design(model, logname="val_monomer_ligand") diff --git a/src/boltzgen/task/predict/data_from_generated.py b/src/boltzgen/task/predict/data_from_generated.py index d282847b..cdb50976 100755 --- a/src/boltzgen/task/predict/data_from_generated.py +++ b/src/boltzgen/task/predict/data_from_generated.py @@ -844,6 +844,8 @@ def transfer_batch_to_device( "tokenized", "data_sample_idx", ]: + if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available(): + batch[key] = batch[key].float() batch[key] = batch[key].to(device) return batch diff --git a/src/boltzgen/task/predict/data_from_yaml.py b/src/boltzgen/task/predict/data_from_yaml.py index 6679da09..c88e96a1 100755 --- a/src/boltzgen/task/predict/data_from_yaml.py +++ b/src/boltzgen/task/predict/data_from_yaml.py @@ -441,5 +441,8 @@ def transfer_batch_to_device( "extra_mols", "data_sample_idx", ]: + #Convert torch.float64 to torch.float for mps compatibility + if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available(): + batch[key] = batch[key].float() batch[key] = batch[key].to(device) return batch diff --git a/src/boltzgen/task/predict/data_ligands.py b/src/boltzgen/task/predict/data_ligands.py index 4b1fe4db..a465c8a9 100755 --- a/src/boltzgen/task/predict/data_ligands.py +++ b/src/boltzgen/task/predict/data_ligands.py @@ -406,5 +406,7 @@ def transfer_batch_to_device( "extra_mols", "data_sample_idx", ]: + if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available(): + batch[key] = batch[key].float() batch[key] = batch[key].to(device) return batch diff --git a/src/boltzgen/task/predict/data_protein_binder.py b/src/boltzgen/task/predict/data_protein_binder.py index 08b88843..2968bbcf 100755 --- a/src/boltzgen/task/predict/data_protein_binder.py +++ b/src/boltzgen/task/predict/data_protein_binder.py @@ -584,6 +584,8 @@ def transfer_batch_to_device( "extra_mols", "data_sample_idx", ]: + if torch.is_tensor(batch[key]) and batch[key].dtype == torch.float64 and torch.backends.mps.is_available(): + batch[key] = batch[key].float() batch[key] = batch[key].to(device) return batch From 0327685c2c1265b17b649b767553d139c1bb3006 Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 31 Mar 2026 15:22:16 +0200 Subject: [PATCH 2/6] Fix MPS compatibility for new upstream code Replace hardcoded torch.autocast("cuda") with device-agnostic device_type=tensor.device.type in confidence_utils, inverse_fold, and writer modules introduced in the upstream merge. Co-Authored-By: Claude Sonnet 4.6 --- src/boltzgen/model/layers/confidence_utils.py | 2 +- src/boltzgen/model/modules/inverse_fold.py | 2 +- src/boltzgen/task/predict/writer.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/boltzgen/model/layers/confidence_utils.py b/src/boltzgen/model/layers/confidence_utils.py index cfaccebe..81c47879 100755 --- a/src/boltzgen/model/layers/confidence_utils.py +++ b/src/boltzgen/model/layers/confidence_utils.py @@ -22,7 +22,7 @@ def compute_frame_pred( resolved_mask=None, inference=False, ): - with torch.amp.autocast("cuda", enabled=False): + with torch.amp.autocast(device_type=pred_atom_coords.device.type, enabled=False): asym_id_token = feats["asym_id"] asym_id_atom = torch.bmm( feats["atom_to_token"].float(), asym_id_token.unsqueeze(-1).float() diff --git a/src/boltzgen/model/modules/inverse_fold.py b/src/boltzgen/model/modules/inverse_fold.py index 56974c12..28d291b5 100755 --- a/src/boltzgen/model/modules/inverse_fold.py +++ b/src/boltzgen/model/modules/inverse_fold.py @@ -470,7 +470,7 @@ def extract_attr_feat( dim=-1, ) - with torch.autocast("cuda", enabled=False): + with torch.autocast(device_type=feats["atom_to_token"].device.type, enabled=False): atom_to_token = feats["atom_to_token"].float() atom_to_token_mean = atom_to_token / ( atom_to_token.sum(dim=1, keepdim=True) + 1e-6 diff --git a/src/boltzgen/task/predict/writer.py b/src/boltzgen/task/predict/writer.py index 551956ec..c792e610 100755 --- a/src/boltzgen/task/predict/writer.py +++ b/src/boltzgen/task/predict/writer.py @@ -416,7 +416,7 @@ def write_on_batch_end( # noqa: PLR0915 traj = trajs[n] aligned = [traj[0]] for frame in traj[1:]: - with torch.autocast("cuda", enabled=False): + with torch.autocast(device_type=frame.device.type, enabled=False): aligned.append( weighted_rigid_align( frame.float().unsqueeze(0), @@ -464,7 +464,7 @@ def write_on_batch_end( # noqa: PLR0915 traj = trajs[n] aligned = [traj[0]] for frame in traj[1:]: - with torch.autocast("cuda", enabled=False): + with torch.autocast(device_type=frame.device.type, enabled=False): aligned.append( weighted_rigid_align( frame.float().unsqueeze(0), From 4bf9161082a6a99864eadabbe4a62d4841ec22b4 Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 31 Mar 2026 18:24:12 +0200 Subject: [PATCH 3/6] Fix KeyError 'name' by reloading molecules fresh from moldir in workers Python pickle does not preserve RDKit atom-level SetProp values. When PyTorch DataLoader spawns worker processes (default num_workers=1 on macOS), self.canonicals is pickled and all atom 'name' properties are lost, causing KeyError in process_atom_features. Fix: load all required molecules directly from the moldir zip inside each get_sample() / get_feat() call instead of using the pickled self.canonicals. The moldir zip handle is cached per-process by _get_zipfile(), so there is no repeated I/O overhead. Co-Authored-By: Claude Sonnet 4.6 --- src/boltzgen/task/predict/data_from_generated.py | 9 +++++---- src/boltzgen/task/predict/data_from_yaml.py | 10 +++++----- src/boltzgen/task/predict/data_ligands.py | 9 ++++----- src/boltzgen/task/predict/data_protein_binder.py | 9 ++++----- 4 files changed, 18 insertions(+), 19 deletions(-) diff --git a/src/boltzgen/task/predict/data_from_generated.py b/src/boltzgen/task/predict/data_from_generated.py index 1ba11de8..c35fbd17 100755 --- a/src/boltzgen/task/predict/data_from_generated.py +++ b/src/boltzgen/task/predict/data_from_generated.py @@ -427,16 +427,17 @@ def get_feat(self, path, design_mask, ss_type=None, binding_type=None, aa_constr try: # Try to find molecules in the dataset moldir if provided # Find missing ones in global moldir and check if all found + # Note: load fresh from moldir to avoid losing RDKit atom properties + # when self.canonicals is pickled by DataLoader worker processes. molecules = {} - molecules.update(self.canonicals) mol_names = set(tokenized.tokens["res_name"].tolist()) - mol_names = mol_names - set(self.canonicals.keys()) if mols is not None: molecules.update(mols) - mol_names = mol_names - set(molecules.keys()) + mol_names = mol_names - set(mols.keys()) if self.moldir is not None: molecules.update(load_molecules(self.moldir, mol_names)) - molecules.update(load_molecules(self.moldir, mol_names)) + else: + molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names}) except Exception as e: # noqa: BLE001 print(f"Molecule loading failed for {path} with error {e}. Skipping.") raise DataFetchException() from e diff --git a/src/boltzgen/task/predict/data_from_yaml.py b/src/boltzgen/task/predict/data_from_yaml.py index f832485e..946a6017 100755 --- a/src/boltzgen/task/predict/data_from_yaml.py +++ b/src/boltzgen/task/predict/data_from_yaml.py @@ -237,16 +237,16 @@ def get_sample(self, path: Path, sample_id: Optional[str] = None) -> Dict: # Try to find molecules in the dataset moldir if provided # Find missing ones in global moldir and check if all found + # Note: self.canonicals may lose atom-level RDKit properties when pickled + # by the DataLoader for worker processes (Python pickle does not preserve + # RDKit atom SetProp values). Load all molecules fresh from moldir instead. molecules = {} - molecules.update(self.canonicals) mol_names = set(tokenized.tokens["res_name"].tolist()) - mol_names = mol_names - set(self.canonicals.keys()) mol_names = mol_names - set(parsed.extra_mols.keys()) if self.moldir is not None: molecules.update(load_molecules(self.moldir, mol_names)) - - mol_names = mol_names - set(molecules.keys()) - molecules.update(load_molecules(self.moldir, mol_names)) + else: + molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names}) molecules.update(parsed.extra_mols) # Finalize input data diff --git a/src/boltzgen/task/predict/data_ligands.py b/src/boltzgen/task/predict/data_ligands.py index a465c8a9..ec22fee1 100755 --- a/src/boltzgen/task/predict/data_ligands.py +++ b/src/boltzgen/task/predict/data_ligands.py @@ -221,15 +221,14 @@ def __getitem__(self, idx: int) -> Dict: try: # Try to find molecules in the dataset moldir if provided # Find missing ones in global moldir and check if all found + # Note: load fresh from moldir to avoid losing RDKit atom properties + # when self.canonicals is pickled by DataLoader worker processes. molecules = {} - molecules.update(self.canonicals) mol_names = set(tokenized.tokens["res_name"].tolist()) - mol_names = mol_names - set(self.canonicals.keys()) if self.moldir is not None: molecules.update(load_molecules(self.moldir, mol_names)) - - mol_names = mol_names - set(molecules.keys()) - molecules.update(load_molecules(self.moldir, mol_names)) + else: + molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names}) except Exception as e: # noqa: BLE001 print(f"Molecule loading failed for {target_id} with error {e}. Skipping.") return self.__getitem__(0) diff --git a/src/boltzgen/task/predict/data_protein_binder.py b/src/boltzgen/task/predict/data_protein_binder.py index 2968bbcf..caad3adb 100755 --- a/src/boltzgen/task/predict/data_protein_binder.py +++ b/src/boltzgen/task/predict/data_protein_binder.py @@ -366,15 +366,14 @@ def __getitem__(self, idx: int) -> Dict: try: # Try to find molecules in the dataset moldir if provided # Find missing ones in global moldir and check if all found + # Note: load fresh from moldir to avoid losing RDKit atom properties + # when self.canonicals is pickled by DataLoader worker processes. molecules = {} - molecules.update(self.canonicals) mol_names = set(tokenized.tokens["res_name"].tolist()) - mol_names = mol_names - set(self.canonicals.keys()) if self.moldir is not None: molecules.update(load_molecules(self.moldir, mol_names)) - - mol_names = mol_names - set(molecules.keys()) - molecules.update(load_molecules(self.moldir, mol_names)) + else: + molecules.update({k: v for k, v in self.canonicals.items() if k in mol_names}) except Exception as e: # noqa: BLE001 print(f"Molecule loading failed for {record.id} with error {e}. Skipping.") return self.__getitem__(0) From aeba138d4a30bf29ce5233db40a3ba88d19e8f87 Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 31 Mar 2026 18:39:19 +0200 Subject: [PATCH 4/6] Fix MPS pin_memory warning and persistent_workers hint in all DataLoaders - Disable pin_memory on MPS (unsupported, causes UserWarning) - Enable persistent_workers when num_workers > 0 (avoids repeated worker init overhead and the PL suggestion warning) Co-Authored-By: Claude Sonnet 4.6 --- src/boltzgen/task/predict/data_from_generated.py | 4 +++- src/boltzgen/task/predict/data_from_yaml.py | 4 +++- src/boltzgen/task/predict/data_ligands.py | 4 +++- src/boltzgen/task/predict/data_protein_binder.py | 4 +++- 4 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/boltzgen/task/predict/data_from_generated.py b/src/boltzgen/task/predict/data_from_generated.py index c35fbd17..a5abedbe 100755 --- a/src/boltzgen/task/predict/data_from_generated.py +++ b/src/boltzgen/task/predict/data_from_generated.py @@ -841,11 +841,13 @@ def output_path_analyzed(input_path): ) def predict_dataloader(self) -> DataLoader: + pin_memory = self.cfg.pin_memory and not torch.backends.mps.is_available() return DataLoader( self.predict_set, batch_size=self.cfg.batch_size, num_workers=self.cfg.num_workers, - pin_memory=self.cfg.pin_memory, + pin_memory=pin_memory, + persistent_workers=self.cfg.num_workers > 0, shuffle=False, collate_fn=collate, ) diff --git a/src/boltzgen/task/predict/data_from_yaml.py b/src/boltzgen/task/predict/data_from_yaml.py index 946a6017..46c3ea61 100755 --- a/src/boltzgen/task/predict/data_from_yaml.py +++ b/src/boltzgen/task/predict/data_from_yaml.py @@ -393,11 +393,13 @@ def predict_dataloader(self) -> DataLoader: The training dataloader. """ + pin_memory = self.pin_memory and not torch.backends.mps.is_available() return DataLoader( self.predict_set, batch_size=self.batch_size, num_workers=self.num_workers, - pin_memory=self.pin_memory, + pin_memory=pin_memory, + persistent_workers=self.num_workers > 0, shuffle=False, collate_fn=collate, ) diff --git a/src/boltzgen/task/predict/data_ligands.py b/src/boltzgen/task/predict/data_ligands.py index ec22fee1..29e7a400 100755 --- a/src/boltzgen/task/predict/data_ligands.py +++ b/src/boltzgen/task/predict/data_ligands.py @@ -351,11 +351,13 @@ def predict_dataloader(self) -> DataLoader: The training dataloader. """ + pin_memory = self.pin_memory and not torch.backends.mps.is_available() return DataLoader( self.predict_set, batch_size=self.batch_size, num_workers=self.num_workers, - pin_memory=self.pin_memory, + pin_memory=pin_memory, + persistent_workers=self.num_workers > 0, shuffle=False, collate_fn=collate, ) diff --git a/src/boltzgen/task/predict/data_protein_binder.py b/src/boltzgen/task/predict/data_protein_binder.py index caad3adb..359b1e54 100755 --- a/src/boltzgen/task/predict/data_protein_binder.py +++ b/src/boltzgen/task/predict/data_protein_binder.py @@ -529,11 +529,13 @@ def predict_dataloader(self) -> DataLoader: The training dataloader. """ + pin_memory = self.pin_memory and not torch.backends.mps.is_available() return DataLoader( self.predict_set, batch_size=self.batch_size, num_workers=self.num_workers, - pin_memory=self.pin_memory, + pin_memory=pin_memory, + persistent_workers=self.num_workers > 0, shuffle=False, collate_fn=collate, ) From 698d09566ae7b5630fc1c8bad27224dce82d8f2c Mon Sep 17 00:00:00 2001 From: fnachon Date: Tue, 31 Mar 2026 18:41:00 +0200 Subject: [PATCH 5/6] Update README.md --- README.md | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/README.md b/README.md index c8c36599..83841599 100755 --- a/README.md +++ b/README.md @@ -106,9 +106,7 @@ docker build -t boltzgen:weights --build-arg DOWNLOAD_WEIGHTS=true . `boltzgen run` takes a [design specification](#how-to-make-a-design-specification-yaml) `.yaml` and produces a set of ranked designs.\ ⚠️ it downloads models (~6GB) to `~/.cache`. This can by changed by passing `--cache YOUR_PATH` or by setting `$HF_HOME`.\ -⚠️ If your run is ever interrupted, you can restart it with `--reuse`. No progress is lost.\ -⚠️ On MacOS set `--num_workers 0` to prevent MPS-related incompatibilities and runtime errors. Optionnaly, to suppress MPS warnings about unsupported pinned memory, set `--config data.pin_memory=false` for the following steps: `design`, `inverse_folding`, `folding`, and `design_folding` - +⚠️ If your run is ever interrupted, you can restart it with `--reuse`. No progress is lost. ```bash boltzgen run example/vanilla_protein/1g13prot.yaml \ From bb23918f65b581a5c2e5cb682088277a45b2ae9e Mon Sep 17 00:00:00 2001 From: fnachon Date: Fri, 10 Jul 2026 17:12:38 +0200 Subject: [PATCH 6/6] Force float32 precision on CPU/MPS instead of bf16-mixed design.yaml, fold.yaml, and affinity.yaml all hardcode trainer.precision: bf16-mixed with no accelerator-conditional override anywhere in the CLI. On CPUs with AVX-512 BF16 support, Lightning's bf16-mixed silently runs ops in bfloat16 rather than falling back, producing structures with wrong bond lengths and atom clashes instead of an error. MPS has the same reliability problem. boltz hit and fixed the identical bug (jwohlwend/boltz#653); this port forces precision=32 whenever trainer.accelerator is explicitly "cpu" or "mps". Co-Authored-By: Claude Sonnet 5 --- src/boltzgen/task/predict/predict.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/boltzgen/task/predict/predict.py b/src/boltzgen/task/predict/predict.py index 1e865f51..e2f348b8 100755 --- a/src/boltzgen/task/predict/predict.py +++ b/src/boltzgen/task/predict/predict.py @@ -114,6 +114,23 @@ def run(self, config: OmegaConf = None, run_prediction=True) -> None: # noqa: A if self.trainer is None: self.trainer = {} + # bf16-mixed silently runs in bfloat16 on CPU (on CPUs with AVX-512 BF16 + # support) and isn't reliably supported on MPS either, producing + # structures with wrong bond lengths and atom clashes instead of an + # error. Force float32 on these accelerators regardless of what the + # step config requested. Same bug as jwohlwend/boltz#653. + accelerator = self.trainer.get("accelerator") + if accelerator in ("cpu", "mps") and self.trainer.get("precision") not in ( + 32, + "32", + "32-true", + ): + print( + f"Accelerator is {accelerator!r}: forcing precision=32 " + "(bf16-mixed silently produces wrong results here)." + ) + self.trainer["precision"] = 32 + # Flip some arguments in debug mode devices = self.trainer.get("devices", 1)