Add LBFGS optimizer as an option

mace/cli/run_train.py

@@ -17,6 +17,7 @@
 import torch.nn.functional
 from e3nn.util import jit
 from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim import LBFGS
 from torch.utils.data import ConcatDataset
 from torch_ema import ExponentialMovingAverage
 
@@ -639,7 +640,7 @@ def run(args) -> None:
             dataset=train_sets[head_config.head_name],
             batch_size=args.batch_size,
             shuffle=True,
-            drop_last=True,
+            drop_last=(not args.lbfgs),
             pin_memory=args.pin_memory,
             num_workers=args.num_workers,
             generator=torch.Generator().manual_seed(args.seed),
@@ -655,7 +656,7 @@ def run(args) -> None:
             num_replicas=world_size,
             rank=rank,
             shuffle=True,
-            drop_last=True,
+            drop_last=(not args.lbfgs),
             seed=args.seed,
         )
         valid_samplers = {}
@@ -674,7 +675,7 @@ def run(args) -> None:
         batch_size=args.batch_size,
         sampler=train_sampler,
         shuffle=(train_sampler is None),
-        drop_last=(train_sampler is None),
+        drop_last=(train_sampler is None and not args.lbfgs),
         pin_memory=args.pin_memory,
         num_workers=args.num_workers,
         generator=torch.Generator().manual_seed(args.seed),
@@ -746,6 +747,8 @@ def run(args) -> None:
     )
 
     start_epoch = 0
+    restart_lbfgs = False
+    opt_start_epoch = None
     if args.restart_latest:
         try:
             opt_start_epoch = checkpoint_handler.load_latest(
@@ -754,11 +757,14 @@ def run(args) -> None:
                 device=device,
             )
         except Exception:  # pylint: disable=W0703
-            opt_start_epoch = checkpoint_handler.load_latest(
-                state=tools.CheckpointState(model, optimizer, lr_scheduler),
-                swa=False,
-                device=device,
-            )
+            try:
+                opt_start_epoch = checkpoint_handler.load_latest(
+                    state=tools.CheckpointState(model, optimizer, lr_scheduler),
+                    swa=False,
+                    device=device,
+                )
+            except Exception: # pylint: disable=W0703
+                restart_lbfgs = True
         if opt_start_epoch is not None:
             start_epoch = opt_start_epoch
 
@@ -769,6 +775,21 @@ def run(args) -> None:
         for group in optimizer.param_groups:
             group["lr"] = args.lr
 
+    if args.lbfgs:
+        logging.info("Switching optimizer to LBFGS")
+        optimizer = LBFGS(model.parameters(),
+                          history_size=200,
+                          max_iter=20,
+                          line_search_fn="strong_wolfe")
+        if restart_lbfgs:
+            opt_start_epoch = checkpoint_handler.load_latest(
+                state=tools.CheckpointState(model, optimizer, lr_scheduler),
+                swa=False,
+                device=device,
+            )
+            if opt_start_epoch is not None:
+                start_epoch = opt_start_epoch
+
     if args.wandb:
         setup_wandb(args)
     if args.distributed:

mace/tools/arg_parser.py

@@ -658,6 +658,12 @@ def build_default_arg_parser() -> argparse.ArgumentParser:
         default=None,
         dest="start_swa",
     )
+    parser.add_argument(
+        "--lbfgs",
+        help="Switch to L-BFGS optimizer",
+        action="store_true",
+        default=False,
+    )
     parser.add_argument(
         "--ema",
         help="use Exponential Moving Average",

mace/tools/train.py

@@ -14,6 +14,7 @@
 import torch
 import torch.distributed
 from torch.nn.parallel import DistributedDataParallel
+from torch.optim import LBFGS
 from torch.optim.swa_utils import SWALR, AveragedModel
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
@@ -182,7 +183,6 @@ def train(
     epoch = start_epoch
 
     # log validation loss before _any_ training
-    valid_loss = 0.0
     for valid_loader_name, valid_loader in valid_loaders.items():
         valid_loss_head, eval_metrics = evaluate(
             model=model,
@@ -230,6 +230,7 @@ def train(
             ema=ema,
             logger=logger,
             device=device,
+            distributed=distributed,
             distributed_model=distributed_model,
             rank=rank,
         )
@@ -247,7 +248,6 @@ def train(
             if "ScheduleFree" in type(optimizer).__name__:
                 optimizer.eval()
             with param_context:
-                valid_loss = 0.0
                 wandb_log_dict = {}
                 for valid_loader_name, valid_loader in valid_loaders.items():
                     valid_loss_head, eval_metrics = evaluate(
@@ -342,25 +342,45 @@ def train_one_epoch(
     ema: Optional[ExponentialMovingAverage],
     logger: MetricsLogger,
     device: torch.device,
+    distributed: bool,
     distributed_model: Optional[DistributedDataParallel] = None,
     rank: Optional[int] = 0,
 ) -> None:
     model_to_train = model if distributed_model is None else distributed_model
-    for batch in data_loader:
-        _, opt_metrics = take_step(
+
+    if isinstance(optimizer, LBFGS):
+        _, opt_metrics = take_step_lbfgs(
             model=model_to_train,
             loss_fn=loss_fn,
-            batch=batch,
+            data_loader=data_loader,
             optimizer=optimizer,
             ema=ema,
             output_args=output_args,
             max_grad_norm=max_grad_norm,
             device=device,
+            distributed=distributed,
+            rank=rank,
         )
         opt_metrics["mode"] = "opt"
         opt_metrics["epoch"] = epoch
         if rank == 0:
             logger.log(opt_metrics)
+    else:
+        for batch in data_loader:
+            _, opt_metrics = take_step(
+                model=model_to_train,
+                loss_fn=loss_fn,
+                batch=batch,
+                optimizer=optimizer,
+                ema=ema,
+                output_args=output_args,
+                max_grad_norm=max_grad_norm,
+                device=device,
+            )
+            opt_metrics["mode"] = "opt"
+            opt_metrics["epoch"] = epoch
+            if rank == 0:
+                logger.log(opt_metrics)
 
 
 def take_step(
@@ -375,19 +395,25 @@ def take_step(
 ) -> Tuple[float, Dict[str, Any]]:
     start_time = time.time()
     batch = batch.to(device)
-    optimizer.zero_grad(set_to_none=True)
     batch_dict = batch.to_dict()
-    output = model(
-        batch_dict,
-        training=True,
-        compute_force=output_args["forces"],
-        compute_virials=output_args["virials"],
-        compute_stress=output_args["stress"],
-    )
-    loss = loss_fn(pred=output, ref=batch)
-    loss.backward()
-    if max_grad_norm is not None:
-        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)
+
+    def closure():
+        optimizer.zero_grad(set_to_none=True)
+        output = model(
+            batch_dict,
+            training=True,
+            compute_force=output_args["forces"],
+            compute_virials=output_args["virials"],
+            compute_stress=output_args["stress"],
+        )
+        loss = loss_fn(pred=output, ref=batch)
+        loss.backward()
+        if max_grad_norm is not None:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)
+
+        return loss
+
+    loss = closure()
     optimizer.step()
 
     if ema is not None:
@@ -401,6 +427,102 @@ def take_step(
     return loss, loss_dict
 
 
+def take_step_lbfgs(
+    model: torch.nn.Module,
+    loss_fn: torch.nn.Module,
+    data_loader: DataLoader,
+    optimizer: torch.optim.Optimizer,
+    ema: Optional[ExponentialMovingAverage],
+    output_args: Dict[str, bool],
+    max_grad_norm: Optional[float],
+    device: torch.device,
+    distributed: bool,
+    rank: int,
+) -> Tuple[float, Dict[str, Any]]:
+    start_time = time.time()
+    logging.debug(
+        f"Max Allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB"
+    )
+
+    total_sample_count = 0
+    for batch in data_loader:
+        total_sample_count += batch.num_graphs
+
+    if distributed:
+        global_sample_count = torch.tensor(total_sample_count, device=device)
+        torch.distributed.all_reduce(
+            global_sample_count, op=torch.distributed.ReduceOp.SUM
+        )
+        total_sample_count = global_sample_count.item()
+
+    signal = torch.zeros(1, device=device) if distributed else None
+
+    def closure():
+        if distributed:
+            if rank == 0:
+                signal.fill_(1)
+                torch.distributed.broadcast(signal, src=0)
+
+            for param in model.parameters():
+                torch.distributed.broadcast(param.data, src=0)
+
+        optimizer.zero_grad(set_to_none=True)
+        total_loss = torch.tensor(0.0, device=device)
+
+        # Process each batch and then collect the results we pass to the optimizer
+        for batch in data_loader:
+            batch = batch.to(device)
+            batch_dict = batch.to_dict()
+            output = model(
+                batch_dict,
+                training=True,
+                compute_force=output_args["forces"],
+                compute_virials=output_args["virials"],
+                compute_stress=output_args["stress"],
+            )
+            batch_loss = loss_fn(pred=output, ref=batch)
+            batch_loss = batch_loss * (batch.num_graphs / total_sample_count)
+
+            batch_loss.backward()
+            total_loss += batch_loss
+
+        if max_grad_norm is not None:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=max_grad_norm)
+
+        if distributed:
+            torch.distributed.all_reduce(total_loss, op=torch.distributed.ReduceOp.SUM)
+        return total_loss
+
+    if distributed:
+        if rank == 0:
+            loss = optimizer.step(closure)
+            signal.fill_(0)
+            torch.distributed.broadcast(signal, src=0)
+        else:
+            while True:
+                # Other ranks wait for signals from rank 0
+                torch.distributed.broadcast(signal, src=0)
+                if signal.item() == 0:
+                    break
+                if signal.item() == 1:
+                    loss = closure()
+
+        for param in model.parameters():
+            torch.distributed.broadcast(param.data, src=0)
+    else:
+        loss = optimizer.step(closure)
+
+    if ema is not None:
+        ema.update()
+
+    loss_dict = {
+        "loss": to_numpy(loss),
+        "time": time.time() - start_time,
+    }
+
+    return loss, loss_dict
+
+
 def evaluate(
     model: torch.nn.Module,
     loss_fn: torch.nn.Module,