diff --git a/README.md b/README.md index 5f6ebd82..83841599 100755 --- a/README.md +++ b/README.md @@ -25,7 +25,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). @@ -36,6 +36,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) @@ -97,7 +108,6 @@ docker build -t boltzgen:weights --build-arg DOWNLOAD_WEIGHTS=true . ⚠️ 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. - ```bash boltzgen run example/vanilla_protein/1g13prot.yaml \ --output workbench/test_run \ diff --git a/pyproject.toml b/pyproject.toml index 6e32f4b1..2a251bec 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 4cf6c6c0..c0d14395 100644 --- a/src/boltzgen/cli/boltzgen.py +++ b/src/boltzgen/cli/boltzgen.py @@ -917,8 +917,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 @@ -937,9 +937,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/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/loss/diffusion.py b/src/boltzgen/model/loss/diffusion.py index 3c192f3e..b11d6c0e 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 6d5ef438..ffb7d1e3 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 @@ -1371,7 +1377,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/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/model/modules/trunk.py b/src/boltzgen/model/modules/trunk.py index fff761c6..e3387eef 100755 --- a/src/boltzgen/model/modules/trunk.py +++ b/src/boltzgen/model/modules/trunk.py @@ -357,7 +357,7 @@ def forward( ).float() # 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) @@ -507,7 +507,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/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()}" diff --git a/src/boltzgen/task/predict/data_from_generated.py b/src/boltzgen/task/predict/data_from_generated.py index 4947e480..a5abedbe 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 @@ -840,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, ) @@ -877,6 +880,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 77fbf84c..46c3ea61 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 @@ -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, ) @@ -445,5 +447,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..29e7a400 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) @@ -352,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, ) @@ -406,5 +407,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..359b1e54 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) @@ -530,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, ) @@ -584,6 +585,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 diff --git a/src/boltzgen/task/predict/writer.py b/src/boltzgen/task/predict/writer.py index 111d30ec..fd93b013 100755 --- a/src/boltzgen/task/predict/writer.py +++ b/src/boltzgen/task/predict/writer.py @@ -415,7 +415,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), @@ -463,7 +463,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),