diff --git a/bin/emle-train b/bin/emle-train index 46136f0..b296161 100755 --- a/bin/emle-train +++ b/bin/emle-train @@ -90,6 +90,23 @@ parser.add_argument( parser.add_argument( "--plot-data", type=str, metavar="name.mat", default=None, help="Data for plotting" ) +parser.add_argument( + "--use-minibatch", + action="store_true", + help="Use minibatch training", +) +parser.add_argument( + "--batch-size", + type=int, + metavar="", + default=1024, + help="Batch size for minibatch training", +) +parser.add_argument( + "--shuffle", + action="store_true", + help="Shuffle training data", +) parser.add_argument("output", type=str, help="Output model file") args = parser.parse_args() @@ -128,4 +145,7 @@ trainer.train( print_every=args.print_every, model_filename=args.output, plot_data_filename=args.plot_data, + use_minibatch=args.use_minibatch, + batch_size=args.batch_size, + shuffle=args.shuffle, ) diff --git a/emle/train/_ivm.py b/emle/train/_ivm.py index ed21b53..d9d4d54 100644 --- a/emle/train/_ivm.py +++ b/emle/train/_ivm.py @@ -24,6 +24,7 @@ import numpy as _np import torch as _torch +from loguru import logger as _logger from ._gpr import GPR as _GPR from ._utils import pad_to_max as _pad_to_max @@ -143,3 +144,124 @@ def perform_ivm(aev_mols, z_mols, atom_ids, species, thr, sigma): aev_ivm_allz = [aev_z[ivm_idx_z] for aev_z, ivm_idx_z in zip(aev_allz, ivm_idx)] return ivm_mol_atom_ids_padded, aev_ivm_allz + + @staticmethod + def perform_ivm_lazy( + aev_filenames, + z_batches, + atom_id_batches, + species, + sigma, + thr, + n_max=None, + device="cuda", + ): + """ + Calculate representative feature vectors for each species from batches loaded from disk. + + Parameters + ---------- + aev_filenames: list of str + Paths to saved AEV batches. + + z_batches: list of torch.Tensor(N_ATOMS_BATCH,) + Atomic numbers for each batch. + + atom_id_batches: list of torch.Tensor(N_ATOMS_BATCH, 2) + (mol_idx, atom_idx) for each atom in each batch. + + species: torch.Tensor(N_SPECIES,) + Unique species to select from. + + sigma: float + Kernel width. + + thr: float + Variance threshold. + + n_max: int or None + Max number of reference AEVs to select per species. + + device: torch.device + Device to use. + + Returns + ------- + ivm_mol_atom_ids_padded: torch.Tensor(N_SPECIES, MAX_N_REF, 2) + IVM selected (mol_idx, atom_idx) per species. + + aev_ivm_allz: list of torch.Tensor(N_REF, AEV_DIM) + AEV features for reference atoms. + """ + selected = {z.item(): [] for z in species} + selected_ids = {z.item(): [] for z in species} + + for z in species: + z = z.item() + iter_count = 0 + while True: + max_var = -float("inf") + best_aev = None + best_id = None + + for fname, z_batch, id_batch in zip( + aev_filenames, z_batches, atom_id_batches + ): + aev_batch = _torch.load(fname, map_location=device).to(device) + z_b = z_batch.to(device) + id_b = id_batch.to(device) + mask = (z_b == z).to(device) + + if not mask.any(): + continue + + aev_z = aev_batch[mask] + ids_z = id_b[mask] + + if len(selected[z]) == 0: + local_var = _torch.ones(len(aev_z), device=device) + else: + aev_sel = _torch.stack(selected[z]).to(device) + k_sel = _GPR._aev_kernel(aev_sel, aev_sel) + k_inv = _torch.linalg.inv( + k_sel + _torch.eye(len(aev_sel), device=device) * sigma**2 + ) + k = _GPR._aev_kernel(aev_z, aev_sel) + local_var = 1 - _torch.sum(k @ k_inv * k, dim=1) + + top_idx = _torch.argmax(local_var) + if local_var[top_idx] > max_var: + max_var = local_var[top_idx].item() + best_aev = aev_z[top_idx].detach().cpu() + best_id = ids_z[top_idx].detach().cpu() + + if max_var < thr or (n_max is not None and len(selected[z]) >= n_max): + break + + selected[z].append(best_aev.to(device)) + selected_ids[z].append(best_id) + iter_count += 1 + _logger.info( + f"IVM for species {z}: Iter {iter_count}: max var = {max_var:.5f}, n_selected = {len(selected[z])}" + ) + + ivm_mol_atom_ids_padded = _pad_to_max( + [ + _torch.stack(v).to(device) + if len(v) > 0 + else _torch.empty(0, 2, dtype=_torch.long, device=device) + for v in selected_ids.values() + ], + value=-1, + ) + + aev_ivm_allz = [ + _torch.stack(v).to(device) + if len(v) > 0 + else _torch.empty( + 0, selected[next(iter(selected))][0].shape[-1], device=device + ) + for v in selected.values() + ] + + return ivm_mol_atom_ids_padded, aev_ivm_allz diff --git a/emle/train/_trainer.py b/emle/train/_trainer.py index 2ab296c..f83547b 100644 --- a/emle/train/_trainer.py +++ b/emle/train/_trainer.py @@ -24,6 +24,8 @@ import os as _os import sys as _sys import torch as _torch +from torch.utils.data import TensorDataset as _TensorDataset +from torch.utils.data import DataLoader as _DataLoader from ..models import EMLEAEVComputer as _EMLEAEVComputer from ..models import EMLEBase as _EMLEBase @@ -170,8 +172,8 @@ def _train_s(s, zid, aev_mols, aev_ivm_allz, sigma): Returns ------- - torch.Tensor(N_BATCH, N_ATOMS) - Atomic widths. + torch.Tensor(N_SPECIES, MAX_N_REF) + Fitted reference values. """ n_ref = _torch.tensor([_.shape[0] for _ in aev_ivm_allz], device=s.device) K_ref_ref_padded, K_mols_ref = _GPR.get_gpr_kernels( @@ -184,6 +186,89 @@ def _train_s(s, zid, aev_mols, aev_ivm_allz, sigma): return _pad_to_max(ref_values_s) + @staticmethod + def _train_s_lazy(s, zid, aev_filenames, aev_ivm_allz, sigma): + """ + Train the s model in a memory-efficient, lazy (batch-wise) manner. + + Parameters + ---------- + s: torch.Tensor(N_BATCH, N_ATOMS) + Atomic widths. + + zid: torch.Tensor(N_BATCH, N_ATOMS) + Species IDs. + + aev_filenames: list of str + List of filenames for AEV batches. + + aev_ivm_allz: list of torch.Tensor(N_REF, AEV_DIM) + Atomic environment vectors for all species (reference AEVs). + + sigma: float + GPR sigma value. + + Returns + ------- + + torch.Tensor(N_SPECIES, MAX_N_REF) + Fitted reference values. + """ + n_ref = _torch.tensor([_.shape[0] for _ in aev_ivm_allz], device=s.device) + n_species = len(aev_ivm_allz) + device = s.device + dtype = s.dtype + + values_by_species = [[] for _ in range(n_species)] + K_mols_ref_by_species = [[] for _ in range(n_species)] + zid_by_species = [[] for _ in range(n_species)] + + offset = 0 + for fname in aev_filenames: + aev_batch = _torch.load(fname, map_location=device) + batch_size = aev_batch.shape[0] + s_batch = s[offset : offset + batch_size] + zid_batch = zid[offset : offset + batch_size] + offset += batch_size + + for i in range(n_species): + mask = zid_batch == i + if not mask.any(): + continue + # aev_z: [N_atoms_of_species, AEV_DIM] + aev_z = aev_batch[mask] + s_z = s_batch[mask] + zid_z = zid_batch[mask] + if aev_z.numel() == 0: + continue + + K_mol_ref = _GPR._aev_kernel(aev_z, aev_ivm_allz[i][: n_ref[i], :]) + values_by_species[i].append(s_z) + K_mols_ref_by_species[i].append(K_mol_ref) + zid_by_species[i].append(zid_z) + + ref_values_s = [] + for i in range(n_species): + if len(values_by_species[i]) == 0: + ref_values_s.append(_torch.zeros(n_ref[i], dtype=dtype, device=device)) + continue + values = _torch.cat(values_by_species[i], dim=0) + K_mols_ref = _torch.cat(K_mols_ref_by_species[i], dim=0) + K_ref_ref = _GPR._aev_kernel( + aev_ivm_allz[i][: n_ref[i], :], aev_ivm_allz[i][: n_ref[i], :] + ) + y_ref = _GPR._fit_sparse_gpr(values, K_mols_ref, K_ref_ref, sigma) + if y_ref.shape[0] < aev_ivm_allz[i].shape[0]: + y_ref_padded = _torch.zeros( + aev_ivm_allz[i].shape[0], dtype=dtype, device=device + ) + y_ref_padded[: y_ref.shape[0]] = y_ref + ref_values_s.append(y_ref_padded) + else: + ref_values_s.append(y_ref) + + return _pad_to_max(ref_values_s) + @staticmethod def _train_model( loss_class, @@ -191,98 +276,110 @@ def _train_model( lr, epochs, emle_base, + loader, print_every=10, + ddp: bool = False, + local_rank: int = 0, *args, **kwargs, ): """ - Train a model. + Train a model using a DataLoader for batching. Optionally uses DistributedDataParallel (DDP) for multi-GPU training. Parameters ---------- - - loss_class: class - Loss class. - - opt_param_names: list of str - List of parameter names to optimize. - - lr: float - Learning rate. - - epochs: int + loss_class : class + The loss class to instantiate for training (e.g., QEqLoss, TholeLoss). + opt_param_names : list of str + List of parameter names to optimize (e.g., ["a_QEq", "ref_values_chi"]). + lr : float + Learning rate for the optimizer. + epochs : int Number of training epochs. - - emle_base: EMLEBase - EMLEBase instance. - - print_every: int - How often to print training progress + emle_base : EMLEBase + An instance of the EMLEBase model. + loader : torch.utils.data.DataLoader + DataLoader yielding batches of training data. Each batch should be a tuple of tensors: + (z, xyz, s, q_core, q_mol, q, alpha, zid), all already on the correct device and dtype. + print_every : int, optional + How often to print training progress (default: 10). + ddp : bool, optional + If True, use DistributedDataParallel for multi-GPU training (default: False). + local_rank : int, optional + Local GPU index for DDP (default: 0). Should be set automatically by torchrun. + *args, **kwargs : + Additional arguments passed to the loss/model call. Returns ------- - - model - Trained model. + model : nn.Module + The trained model (loss instance). """ + model = loss_class(emle_base).to( + f"cuda:{local_rank}" if ddp else emle_base._device + ) - def _train_loop( - loss_instance, optimizer, epochs, print_every=10, *args, **kwargs - ): - """ - Perform the training loop. - - Parameters - ---------- - - loss_instance: nn.Module - Loss instance. - - optimizer: torch.optim.Optimizer - Optimizer. - - epochs: int - Number of training epochs. - - print_every: int - How often to print training progress - - args: list - Positional arguments to pass to the forward method. - - kwargs: dict - Keyword arguments to pass to the forward method. - - Returns - ------- - - loss - Forward loss. - """ - for epoch in range(epochs): - loss_instance.train() - optimizer.zero_grad() - loss, rmse, max_error = loss_instance(*args, **kwargs) - loss.backward(retain_graph=True) - optimizer.step() - if (epoch + 1) % print_every == 0: - _logger.info( - f"Epoch {epoch+1}: Loss ={loss.item():9.4f} " - f"RMSE ={rmse.item():9.4f} " - f"Max Error ={max_error.item():9.4f}" - ) + if ddp: + import torch.distributed as dist + from torch.nn.parallel import DistributedDataParallel as DDP - return loss + model = DDP(model, device_ids=[local_rank]) - model = loss_class(emle_base) + loss_name = loss_class.__name__.lower() opt_parameters = [ param for name, param in model.named_parameters() if name.split(".")[1] in opt_param_names ] - optimizer = _torch.optim.Adam(opt_parameters, lr=lr) - _train_loop(model, optimizer, epochs, print_every, *args, **kwargs) + + for epoch in range(epochs): + if ddp and hasattr(loader.sampler, "set_epoch"): + loader.sampler.set_epoch(epoch) + model.train() + running_loss = 0.0 + running_sq_error = 0.0 + running_count = 0 + running_max_error = 0.0 + for batch in loader: + z_b, xyz_b, s_b, q_core_b, q_mol_b, q_b, alpha_b, zid_b = batch + optimizer.zero_grad() + if loss_name == "qeqloss": + loss, rmse, max_error = model( + atomic_numbers=z_b, + xyz=xyz_b, + q_mol=q_mol_b, + q_target=q_b, + **kwargs, + ) + elif loss_name == "tholeloss": + loss, rmse, max_error = model( + atomic_numbers=z_b, + xyz=xyz_b, + q_mol=q_mol_b, + alpha_mol_target=alpha_b, + **kwargs, + ) + else: + raise ValueError(f"Unsupported loss class: {loss_name}") + loss.backward(retain_graph=True) + optimizer.step() + batch_size_actual = z_b.shape[0] + running_loss += loss.item() * batch_size_actual + running_sq_error += (rmse.item() ** 2) * batch_size_actual + running_count += batch_size_actual + running_max_error = max(running_max_error, max_error.item()) + epoch_loss = running_loss / running_count + epoch_rmse = ( + (running_sq_error / running_count) ** 0.5 if running_count > 0 else 0.0 + ) + epoch_max_error = running_max_error + if (epoch + 1) % print_every == 0 and (not ddp or local_rank == 0): + _logger.info( + f"Epoch {epoch + 1}: Loss ={epoch_loss:9.4f} " + f"RMSE ={epoch_rmse:9.4f} " + f"Max Error ={epoch_max_error:9.4f}" + ) return model def train( @@ -307,10 +404,16 @@ def train( model_filename="emle_model.mat", plot_data_filename=None, device=_torch.device("cuda"), - dtype=_torch.float64, + dtype=_torch.float32, + use_minibatch=False, + batch_size=100, + shuffle=False, + ddp: bool = False, + local_rank: int = 0, + ivm_data_file: str = None, ): """ - Train an EMLE model. + Train an EMLE model, optionally using DistributedDataParallel (DDP) for multi-GPU training. Parameters ---------- @@ -324,12 +427,7 @@ def train( s: numpy.array, List[numpy.array], torch.Tensor, List[torch.Tensor] (N_BATCH, N_ATOMS) Atomic widths. - q_core: numpy.array, List[numpy.array], torch.Tensor, List[torch.Tensor] (N_BATCH, N_ATOMS) - Atomic core charges. - - q_val: array or tensor or list of tensor/arrays of shape (N_BATCH, N_ATOMS) - Atomic valence charges. - + q_core: numpy.array, List[numpy.arrayTrue alpha: array or tensor or list of tensor/arrays of shape (N_BATCH, 3, 3) Atomic polarizabilities. @@ -378,12 +476,40 @@ def train( dtype: torch.dtype Data type to use for training. Default is torch.float64. + use_minibatch: bool + Use minibatch training. Default is False. + + batch_size: int + Batch size for minibatch training. Default is 1024. + + shuffle: bool + Shuffle training data. Default is False. + + ddp : bool, optional + If True, use DistributedDataParallel for multi-GPU training (default: False). + + local_rank : int, optional + Local GPU index for DDP (default: 0). Should be set automatically by torchrun. + + ivm_data_file: str or None + Filename to save IVM data. If None, IVM data is not saved. + Returns ------- dict Trained EMLE model. """ + if ddp: + import torch.distributed as dist + from torch.utils.data.distributed import ( + DistributedSampler as _DistributedSampler, + ) + + dist.init_process_group(backend="nccl") + _torch.cuda.set_device(local_rank) + device = _torch.device(f"cuda:{local_rank}") + # Check input data. assert ( len(z) == len(xyz) == len(s) == len(q_core) == len(q_val) == len(alpha) @@ -406,45 +532,80 @@ def train( q_mol = _torch.sum(q, dim=1) z = _pad_to_max(z) xyz = _pad_to_max(xyz) + s = _pad_to_max(s) + alpha = _pad_to_max(alpha) + + print("Size of the dataset ", len(z)) + # Apply train_mask q_core_train = q_core[train_mask] q_mol_train = q_mol[train_mask] q_train = q[train_mask] z_train = z[train_mask] xyz_train = xyz[train_mask] - s_train = _pad_to_max(s)[train_mask] - alpha_train = _pad_to_max(alpha)[train_mask] - species = _torch.unique(_torch.tensor(z_train[z_train > 0], device=device)) + s_train = s[train_mask] + alpha_train = alpha[train_mask] - # Place on the correct device and set the data type. - q_mol = q_mol.to(device=device, dtype=dtype) + q_core_train = q_core_train.to(device=device, dtype=dtype) q_mol_train = q_mol_train.to(device=device, dtype=dtype) + q_train = q_train.to(device=device, dtype=dtype) z_train = z_train.to(device=device, dtype=_torch.int64) xyz_train = xyz_train.to(device=device, dtype=dtype) s_train = s_train.to(device=device, dtype=dtype) - q_core_train = q_core_train.to(device=device, dtype=dtype) - q_train = q_train.to(device=device, dtype=dtype) alpha_train = alpha_train.to(device=device, dtype=dtype) + + del q_core, q_val, q, q_mol, z, xyz, s, alpha + _torch.cuda.empty_cache() + + # Get unique species + species = _torch.unique(z_train[z_train > 0]) species = species.to(device=device, dtype=_torch.int64) - # Get zid mapping. + # Get zid mapping zid_mapping = self._get_zid_mapping(species) zid_train = zid_mapping[z_train] if computer_n_species is None: computer_n_species = len(species) - # Calculate AEVs. + batch_size_eff = len(z_train) if not use_minibatch else batch_size + dataset = _TensorDataset( + z_train, + xyz_train, + s_train, + q_core_train, + q_mol_train, + q_train, + alpha_train, + zid_train, + ) + if ddp: + sampler = _DistributedSampler(dataset) + loader = _DataLoader(dataset, batch_size=batch_size_eff, sampler=sampler) + else: + loader = _DataLoader(dataset, batch_size=batch_size_eff, shuffle=shuffle) + + # Calculate AEV mask globally emle_aev_computer = _EMLEAEVComputer( num_species=computer_n_species, zid_map=computer_zid_map, dtype=dtype, device=device, ) - aev_mols = emle_aev_computer(zid_train, xyz_train) - aev_mask = _torch.sum(aev_mols.reshape(-1, aev_mols.shape[-1]) ** 2, dim=0) > 0 + aev_mask = None + for batch in loader: + z_b, xyz_b, *_, zid_b = batch + aev_batch = emle_aev_computer(zid_b, xyz_b) + batch_mask = ( + _torch.sum(aev_batch.reshape(-1, aev_batch.shape[-1]) ** 2, dim=0) > 0 + ) + if aev_mask is None: + aev_mask = batch_mask + else: + aev_mask |= batch_mask + del aev_batch + _torch.cuda.empty_cache() - aev_mols = aev_mols[:, :, aev_mask] emle_aev_computer = _EMLEAEVComputer( num_species=computer_n_species, zid_map=computer_zid_map, @@ -453,22 +614,78 @@ def train( device=device, ) + # Save masked AEVs to files if using lazy mode, else keep in memory + aev_filenames = [] + if use_minibatch: + batch_dir = "batches" + _os.makedirs(batch_dir, exist_ok=True) + for i, batch in enumerate(loader): + if not ddp or local_rank == 0: + _logger.info(f"Saving masked AEVs for batch {i + 1}/{len(loader)}") + z_b, xyz_b, *_, zid_b = batch + aev_batch = emle_aev_computer(zid_b, xyz_b) + filename = _os.path.join(batch_dir, f"aev_mols_batch_{i}.pt") + _torch.save(aev_batch.cpu(), filename) + aev_filenames.append(filename) + del aev_batch + _torch.cuda.empty_cache() + else: + aev_mols = emle_aev_computer(zid_mapping[z_train], xyz_train) + # "Fit" q_core (just take averages over the entire training set). q_core_z = _mean_by_z(q_core_train, zid_train) - _logger.info("Performing IVM...") - # Create an array of (molecule_id, atom_id) pairs (as in the full - # dataset) for the training set. This is needed to be able to locate - # atoms/molecules in the original dataset that were picked by IVM. - n_mols, max_atoms = q_train.shape - atom_ids = _torch.stack( - _torch.meshgrid(_torch.arange(n_mols), _torch.arange(max_atoms)), dim=-1 - ).to(device) - - # Perform IVM. - ivm_mol_atom_ids_padded, aev_ivm_allz = _IVM.perform_ivm( - aev_mols, z_train, atom_ids, species, ivm_thr, sigma - ) + # IVM selection + if ivm_data_file is not None and _os.path.exists(ivm_data_file): + if not ddp or local_rank == 0: + _logger.info(f"Loading IVM data from file: {ivm_data_file}") + ivm_data = _torch.load(ivm_data_file, map_location=device) + ivm_mol_atom_ids_padded = ivm_data["ivm_mol_atom_ids_padded"] + aev_ivm_allz = ivm_data["aev_ivm_allz"] + else: + if not use_minibatch: + _logger.info("Performing IVM...") + n_mols, max_atoms = q_train.shape + atom_ids = _torch.stack( + _torch.meshgrid(_torch.arange(n_mols), _torch.arange(max_atoms)), + dim=-1, + ).to(device) + ivm_mol_atom_ids_padded, aev_ivm_allz = _IVM.perform_ivm( + aev_mols, z_train, atom_ids, species, ivm_thr, sigma + ) + else: + _logger.info("Performing Lazy IVM...") + n_mols, max_atoms = q_train.shape + atom_ids = _torch.arange(max_atoms) + z_mols_batches = [] + atom_ids_batches = [] + for i, batch in enumerate(loader): + z_b = batch[0] + batch_size_actual = z_b.shape[0] + start = i * loader.batch_size + mol_range = _torch.arange(start, start + batch_size_actual) + atom_grid = _torch.stack( + _torch.meshgrid(mol_range, atom_ids, indexing="ij"), dim=-1 + ) + z_mols_batches.append(z_b) + atom_ids_batches.append(atom_grid) + ivm_mol_atom_ids_padded, aev_ivm_allz = _IVM.perform_ivm_lazy( + aev_filenames=aev_filenames, + z_batches=z_mols_batches, + atom_id_batches=atom_ids_batches, + species=species, + thr=ivm_thr, + sigma=sigma, + ) + if not ddp or local_rank == 0: + _logger.info(f"Saving IVM data to file: {ivm_data_file}") + _torch.save( + { + "ivm_mol_atom_ids_padded": ivm_mol_atom_ids_padded, + "aev_ivm_allz": aev_ivm_allz, + }, + ivm_data_file, + ) ref_features = _pad_to_max(aev_ivm_allz) ref_mask = ivm_mol_atom_ids_padded[:, :, 0] > -1 @@ -478,7 +695,14 @@ def train( _logger.info(f"{atom_z:2d}: {n:5d}") # Fit s (pure GPR, no fancy optimization needed). - ref_values_s = self._train_s(s_train, zid_train, aev_mols, aev_ivm_allz, sigma) + if not use_minibatch: + ref_values_s = self._train_s( + s_train, zid_train, aev_mols, aev_ivm_allz, sigma + ) + else: + ref_values_s = self._train_s_lazy( + s_train, zid_train, aev_filenames, aev_ivm_allz, sigma + ) # Good for debugging # _torch.autograd.set_detect_anomaly(True) @@ -520,7 +744,8 @@ def train( ) # Fit chi, a_QEq (QEq over chi predicted with GPR). - _logger.info("Fitting a_QEq and chi values...") + if not ddp or local_rank == 0: + _logger.info("Fitting a_QEq and chi values...") self._train_model( loss_class=self._qeq_loss, opt_param_names=["a_QEq", "ref_values_chi"], @@ -528,19 +753,17 @@ def train( epochs=epochs, print_every=print_every, emle_base=emle_base, - atomic_numbers=z_train, - xyz=xyz_train, - q_mol=q_mol_train, - q_target=q_train, + loader=loader, + ddp=ddp, + local_rank=local_rank, ) - # Update GPR constants for chi - # (now inconsistent since not updated after the last epoch) self._qeq_loss._update_chi_gpr(emle_base) - - _logger.debug(f"Optimized a_QEq: {emle_base.a_QEq.data.item()}") + if not ddp or local_rank == 0: + _logger.debug(f"Optimized a_QEq: {emle_base.a_QEq.data.item()}") # Fit a_Thole, k_Z (uses volumes predicted by QEq model). - _logger.info("Fitting a_Thole and k_Z values...") + if not ddp or local_rank == 0: + _logger.info("Fitting a_Thole and k_Z values...") self._train_model( loss_class=self._thole_loss, opt_param_names=["a_Thole", "k_Z"], @@ -548,16 +771,17 @@ def train( epochs=epochs, print_every=print_every, emle_base=emle_base, - atomic_numbers=z_train, - xyz=xyz_train, - q_mol=q_mol_train, - alpha_mol_target=alpha_train, + loader=loader, + ddp=ddp, + local_rank=local_rank, ) + if not ddp or local_rank == 0: + _logger.debug(f"Optimized a_Thole: {emle_base.a_Thole.data.item()}") - _logger.debug(f"Optimized a_Thole: {emle_base.a_Thole.data.item()}") # Fit sqrtk_ref ( alpha = sqrtk ** 2 * k_Z * v). if alpha_mode == "reference": - _logger.info("Fitting ref_values_sqrtk values...") + if not ddp or local_rank == 0: + _logger.info("Fitting ref_values_sqrtk values...") self._train_model( loss_class=self._thole_loss, opt_param_names=["ref_values_sqrtk"], @@ -565,41 +789,40 @@ def train( epochs=epochs, print_every=print_every, emle_base=emle_base, - atomic_numbers=z_train, - xyz=xyz_train, - q_mol=q_mol_train, - alpha_mol_target=alpha_train, + loader=loader, + ddp=ddp, + local_rank=local_rank, opt_sqrtk=True, l2_reg=20.0, ) - # Update GPR constants for sqrtk - # (now inconsistent since not updated after the last epoch) self._thole_loss._update_sqrtk_gpr(emle_base) - # Create the final model. - emle_model = { - "q_core": q_core_z, - "a_QEq": emle_base.a_QEq, - "a_Thole": emle_base.a_Thole, - "s_ref": emle_base.ref_values_s, - "chi_ref": emle_base.ref_values_chi, - "k_Z": emle_base.k_Z, - "sqrtk_ref": ( - emle_base.ref_values_sqrtk if alpha_mode == "reference" else None - ), - "species": species, - "alpha_mode": alpha_mode, - "n_ref": n_ref, - "ref_aev": ref_features, - "aev_mask": aev_mask, - "zid_map": emle_aev_computer._zid_map, - "computer_n_species": computer_n_species, - } - - if model_filename is not None: + # Only save model on rank 0 + if (model_filename is not None) and (not ddp or local_rank == 0): + emle_model = { + "q_core": q_core_z, + "a_QEq": emle_base.a_QEq, + "a_Thole": emle_base.a_Thole, + "s_ref": emle_base.ref_values_s, + "chi_ref": emle_base.ref_values_chi, + "k_Z": emle_base.k_Z, + "sqrtk_ref": ( + emle_base.ref_values_sqrtk if alpha_mode == "reference" else None + ), + "species": species, + "alpha_mode": alpha_mode, + "n_ref": n_ref, + "ref_aev": ref_features, + "aev_mask": aev_mask, + "zid_map": emle_aev_computer._zid_map, + "computer_n_species": computer_n_species, + } self._write_model_to_file(emle_model, model_filename) if plot_data_filename is None: + if ddp: + dist.barrier() + dist.destroy_process_group() return emle_base emle_base._alpha_mode = "species" @@ -632,6 +855,10 @@ def train( A_thole, z_mask ) - self._write_model_to_file(plot_data, plot_data_filename) + if not ddp or local_rank == 0: + self._write_model_to_file(plot_data, plot_data_filename) + if ddp: + dist.barrier() + dist.destroy_process_group() return emle_base