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 ec48b09ed7cc026f3798302369d881f6668a3e36 Mon Sep 17 00:00:00 2001 From: fnachon Date: Fri, 10 Jul 2026 16:50:51 +0200 Subject: [PATCH 6/6] Fix num_filters_passed to count individually-passing filters The counter incremented by the cumulative AND of all filter columns seen so far in the loop, not by whether the current filter passed. Once any filter failed, every later increment added 0 regardless of whether the design actually passed those later filters, so the column ends up reporting "count of leading filters passed before the first failure" instead of the intended total pass count. Fixes HannesStark/boltzgen#126 Co-Authored-By: Claude Sonnet 5 --- src/boltzgen/task/filter/filter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/boltzgen/task/filter/filter.py b/src/boltzgen/task/filter/filter.py index bbb31b2e..44bc0d86 100644 --- a/src/boltzgen/task/filter/filter.py +++ b/src/boltzgen/task/filter/filter.py @@ -408,7 +408,7 @@ def filter_df(self): else: self.df[filter_col] = self.df[feat] >= threshold - self.df["num_filters_passed"] += self.df[filter_cols].all(axis=1) + self.df["num_filters_passed"] += self.df[filter_col].astype(int) self.df["pass_filters"] = self.df[filter_cols].all(axis=1) msg = f"Num designs that pass the {feat} filter with threshold {threshold} where {'lower' if low else 'higher'} is better: {self.df[filter_col].sum()}"