update maximum loss for amplify recipe

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

recipes/esm2_accelerate/tests/test_accelerate_amplify.py

@@ -26,37 +26,37 @@
 
 def test_te_with_default_config(tmp_path):
     train_loss = launch_accelerate("default.yaml", tmp_path, 1, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
 def test_te_with_dynamo_config(tmp_path):
     train_loss = launch_accelerate("dynamo.yaml", tmp_path, 1, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
 def test_te_with_fp8_config(tmp_path):
     train_loss = launch_accelerate("fp8.yaml", tmp_path, 1, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
 def test_te_with_fsdp2_config(tmp_path):
     train_loss = launch_accelerate("fsdp2_te.yaml", tmp_path, 1, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
 @requires_multi_gpu
 def test_te_with_default_config_two_gpus(tmp_path):
     train_loss = launch_accelerate("default.yaml", tmp_path, 2, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
 @requires_multi_gpu
 def test_te_with_fp8_config_two_gpus(tmp_path):
     train_loss = launch_accelerate("fp8.yaml", tmp_path, 2, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"
 
 
 @requires_multi_gpu
 def test_te_with_fsdp2_config_two_gpus(tmp_path):
     train_loss = launch_accelerate("fsdp2_te.yaml", tmp_path, 2, "L0_sanity_amplify")
-    assert train_loss < 6.0, f"Final train_loss {train_loss} should be less than 6.0"
+    assert train_loss < 3.0, f"Final train_loss {train_loss} should be less than 3.0"

recipes/esm2_accelerate/train.py

@@ -47,7 +47,6 @@ def main(args: DictConfig):
 
     config = AutoConfig.from_pretrained(args.model_tag, trust_remote_code=True)
     model = AutoModelForMaskedLM.from_config(config, trust_remote_code=True, dtype=torch.bfloat16)
-    model.accepts_loss_kwargs = False
 
     train_dataset, eval_dataset, data_collator = create_datasets_and_collator(
         tokenizer_name=args.model_tag,