Update for bf16 and fp32

linkedin · Feb 15, 2025 · 762165b · 762165b
1 parent 02dbb2e
commit 762165b
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 7 deletions.
diff --git a/test/convergence/bf16/test_mini_models_multimodal.py b/test/convergence/bf16/test_mini_models_multimodal.py
@@ -322,6 +322,7 @@ def run_mini_model_multimodal(
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
         output = model(**batch)
+        output.logits.retain_grad()
         output.loss.backward()
         optimizer.step()
 
@@ -405,8 +406,8 @@ def test_mini_model_multimodal(
 
     # Compare the logits from the last step
     assert_verbose_allclose(
-        expected_output["logits"],
-        actual_output["logits"],
+        expected_output["logits"].grad,
+        actual_output["logits"].grad,
         atol=logits_atol,
         rtol=logits_rtol,
     )

diff --git a/test/convergence/bf16/test_mini_models_with_logits.py b/test/convergence/bf16/test_mini_models_with_logits.py
@@ -444,7 +444,7 @@ def run_mini_model(
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
         output = model(**batch)
-        output.logits.retain_grad()  # For comparing logits.grad
+        output.logits.retain_grad()
         output.loss.backward()
         optimizer.step()
         print(f"Step {i}, Loss: {output.loss.item()}")

diff --git a/test/convergence/fp32/test_mini_models_multimodal.py b/test/convergence/fp32/test_mini_models_multimodal.py
@@ -321,6 +321,7 @@ def run_mini_model_multimodal(
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
         output = model(**batch)
+        output.logits.retain_grad()
         output.loss.backward()
         optimizer.step()
 
@@ -400,8 +401,8 @@ def test_mini_model_multimodal(
 
     # Compare the logits from the last step
     assert_verbose_allclose(
-        expected_output["logits"],
-        actual_output["logits"],
+        expected_output["logits"].grad,
+        actual_output["logits"].grad,
         atol=logits_atol,
         rtol=logits_rtol,
     )

diff --git a/test/convergence/fp32/test_mini_models_with_logits.py b/test/convergence/fp32/test_mini_models_with_logits.py
@@ -443,6 +443,7 @@ def run_mini_model(
         batch = next(loader_iter).to(model.device)
         optimizer.zero_grad()
         output = model(**batch)
+        output.logits.retain_grad()
         output.loss.backward()
         optimizer.step()
         print(f"Step {i}, Loss: {output.loss.item()}")
@@ -529,8 +530,8 @@ def test_mini_model(
     # import pdb; pdb.set_trace()
     # Compare the logits from the last step
     assert_verbose_allclose(
-        expected_output["logits"],
-        actual_output["logits"],
+        expected_output["logits"].grad,
+        actual_output["logits"].grad,
         atol=logits_atol,
         rtol=logits_rtol,
     )