Fix daily test failures (#3733)

tests/algorithms/test_algorithm_resumption.py

@@ -53,7 +53,7 @@ def test_algorithm_resumption(
 
         copied_model = copy.deepcopy(model)  # copy the model so the params will start from the same point
 
-        optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
         scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
 
         # Reduce training duration and data
@@ -82,7 +82,7 @@ def test_algorithm_resumption(
         # create second trainer, load an intermediate checkpoint
         # and continue training
 
-        optimizer = torch.optim.SGD(copied_model.parameters(), lr=0.1)
+        optimizer = torch.optim.Adam(copied_model.parameters(), lr=0.1)
         scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
 
         alg = alg_cls(**alg_kwargs)

tests/test_full_nlp.py

@@ -60,7 +60,7 @@ def pretraining_test_helper(tokenizer, model, algorithms, tmp_path, device):
         model=pretraining_model_copy,
         train_dataloader=pretraining_train_dataloader,
         save_folder=str(tmp_path / 'pretraining_checkpoints'),
-        max_duration='1ba',
+        max_duration='2ba',
         seed=17,
         algorithms=algorithms,
         device=device,
@@ -138,7 +138,7 @@ def finetuning_test_helper(tokenizer, model, algorithms, checkpoint_path, pretra
         load_weights_only=True,
         load_strict_model_weights=False,
         loggers=[rud],
-        max_duration='1ba',
+        max_duration='2ba',
         seed=17,
         algorithms=algorithms,
         device=device,