From 10fdac4759b6607901b375ecf39c01dd073ad798 Mon Sep 17 00:00:00 2001 From: Gianluca Martino Date: Thu, 23 Jul 2026 18:21:02 -0700 Subject: [PATCH 1/2] Make batched vs list-model training comparison platform-robust --- .../bayesian/test_model_constructor.py | 58 ++++++++++++++----- 1 file changed, 43 insertions(+), 15 deletions(-) diff --git a/xopt/tests/generators/bayesian/test_model_constructor.py b/xopt/tests/generators/bayesian/test_model_constructor.py index b3f043d7..9405ff52 100644 --- a/xopt/tests/generators/bayesian/test_model_constructor.py +++ b/xopt/tests/generators/bayesian/test_model_constructor.py @@ -995,7 +995,7 @@ def test_train_model_batch_compare(self): for i in range(test_vocs.n_outputs): assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=5e-3) - def test_train_model_batch_compare_adam(self): + def test_train_model_batch_compare_training(self): test_vocs = deepcopy(TEST_VOCS) test_data = deepcopy(TEST_DATA) @@ -1020,17 +1020,48 @@ def test_train_model_batch_compare_adam(self): batch_ls = model_single.covar_module.raw_lengthscale for i in range(test_vocs.n_outputs): assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=1e-3) - print([(p, p.shape) for p in model_list.parameters()]) - print("------------------------------") - print([(p, p.shape) for p in model_single.parameters()]) - optimizer_single = torch.optim.Adam(model_single.parameters(), lr=0.1) mll = ExactMarginalLogLikelihood(model_single.likelihood, model_single) - optimizer_list = [] - mll_list = [] - for ml in model_list.models: - optimizer_list.append(torch.optim.Adam(ml.parameters(), lr=0.1)) - mll_list.append(ExactMarginalLogLikelihood(ml.likelihood, ml)) + mll_list = [ + ExactMarginalLogLikelihood(ml.likelihood, ml) for ml in model_list.models + ] + + # The batched MLL decouples exactly across outputs, so before any + # optimizer step the loss and gradients must match tightly + model_single.zero_grad() + loss = -mll( + model_single(model_single.train_inputs[0]), model_single.train_targets + ).sum() + loss.backward() + single_loss = float(loss.item()) + + total_list_loss = 0.0 + for j, ml in enumerate(model_list.models): + ml.zero_grad() + list_loss = -mll_list[j](ml(ml.train_inputs[0]), ml.train_targets) + list_loss.backward() + total_list_loss += list_loss.item() + assert np.isclose(total_list_loss, single_loss, rtol=0.0, atol=1e-9) + + batch_params = dict(model_single.named_parameters()) + for j, ml in enumerate(model_list.models): + for name, p_list in ml.named_parameters(): + grad_batch = batch_params[name].grad + if grad_batch.shape != p_list.grad.shape: + assert grad_batch.shape[0] == test_vocs.n_outputs, name + grad_batch = grad_batch[j] + assert torch.allclose(grad_batch, p_list.grad, rtol=0, atol=1e-9), name + + # Compare training trajectories with SGD: unlike Adam, whose per-step + # normalization chaotically amplifies floating-point reduction-order + # differences between batched and per-model kernels, plain SGD stays + # at float64 noise level for this step/lr budget (larger budgets can + # leave the well-conditioned region and go NotPSD - recheck tolerances + # before increasing them) + optimizer_single = torch.optim.SGD(model_single.parameters(), lr=0.1) + optimizer_list = [ + torch.optim.SGD(ml.parameters(), lr=0.1) for ml in model_list.models + ] for i in range(10): optimizer_single.zero_grad() @@ -1039,7 +1070,6 @@ def test_train_model_batch_compare_adam(self): loss.backward() optimizer_single.step() single_loss = float(loss.item()) - print(f"Single: {loss.item()}") total_list_loss = 0.0 for j, ml in enumerate(model_list.models): optimizer_list[j].zero_grad() @@ -1047,17 +1077,15 @@ def test_train_model_batch_compare_adam(self): loss = -mll_list[j](output, ml.train_targets) loss.backward() optimizer_list[j].step() - print(f"List {j}: {loss.item()}") total_list_loss += loss.item() - print(f"List total: {total_list_loss}") - assert np.isclose(total_list_loss, single_loss, rtol=0.0, atol=1e-5) + assert np.isclose(total_list_loss, single_loss, rtol=0.0, atol=1e-8) list_ls = [ model_list.models[i].covar_module.raw_lengthscale for i in range(test_vocs.n_outputs) ] batch_ls = model_single.covar_module.raw_lengthscale for i in range(test_vocs.n_outputs): - assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=1e-3) + assert torch.allclose(list_ls[i], batch_ls[i, ...], rtol=0, atol=1e-8) @pytest.mark.parametrize("use_cuda", cuda_combinations) def test_batched_basic(self, use_cuda): From 0207bd5665d2101d75c99a0dac0380623bb16928 Mon Sep 17 00:00:00 2001 From: nikitakuklev Date: Fri, 24 Jul 2026 19:59:48 -0500 Subject: [PATCH 2/2] Clarify batched training test comments --- .../bayesian/test_model_constructor.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/xopt/tests/generators/bayesian/test_model_constructor.py b/xopt/tests/generators/bayesian/test_model_constructor.py index 9405ff52..e4caed16 100644 --- a/xopt/tests/generators/bayesian/test_model_constructor.py +++ b/xopt/tests/generators/bayesian/test_model_constructor.py @@ -1026,8 +1026,9 @@ def test_train_model_batch_compare_training(self): ExactMarginalLogLikelihood(ml.likelihood, ml) for ml in model_list.models ] - # The batched MLL decouples exactly across outputs, so before any - # optimizer step the loss and gradients must match tightly + # For independent outputs, the batched MLL is mathematically the sum of + # the per-model MLLs. Compare losses and parameter gradients before + # optimizer state can amplify numerical differences. model_single.zero_grad() loss = -mll( model_single(model_single.train_inputs[0]), model_single.train_targets @@ -1052,12 +1053,11 @@ def test_train_model_batch_compare_training(self): grad_batch = grad_batch[j] assert torch.allclose(grad_batch, p_list.grad, rtol=0, atol=1e-9), name - # Compare training trajectories with SGD: unlike Adam, whose per-step - # normalization chaotically amplifies floating-point reduction-order - # differences between batched and per-model kernels, plain SGD stays - # at float64 noise level for this step/lr budget (larger budgets can - # leave the well-conditioned region and go NotPSD - recheck tolerances - # before increasing them) + # Use SGD for the trajectory check so batched-vs-per-model rounding + # differences remain proportional to the gradient differences. Adam can + # amplify near-zero differences through its adaptive denominator. + # Keep the step count and learning rate small to avoid ill-conditioned + # covariance matrices. optimizer_single = torch.optim.SGD(model_single.parameters(), lr=0.1) optimizer_list = [ torch.optim.SGD(ml.parameters(), lr=0.1) for ml in model_list.models