Refactor dpo and ppo recipe by introducing disable_dropout utility function

Author: Ankur-singh

recipes/full_dpo_distributed.py

@@ -20,7 +20,7 @@
 from torchtune.data import CROSS_ENTROPY_IGNORE_IDX, padded_collate_dpo
 from torchtune.datasets import ConcatDataset
 from torchtune.recipe_interfaces import FTRecipeInterface
-from torchtune.training import DummyProfiler, PROFILER_KEY
+from torchtune.training import disable_dropout, DummyProfiler, PROFILER_KEY
 from torchtune.training.lr_schedulers import get_lr
 from torchtune.utils import get_world_size_and_rank
 from tqdm import tqdm
@@ -494,12 +494,7 @@ def _setup_model(
 
         # disabling dropout if found - non-determinism leads to issues in e.g. comparing logprobs
         # between ref policy and current policy
-        for module in model.modules():
-            if isinstance(module, torch.nn.Dropout):
-                warn(
-                    f"Dropout found in {module}. This is likely to cause issues during training. Disabling."
-                )
-                module.p = 0
+        disable_dropout(model)
 
         # synchronize before training begins
         torch.distributed.barrier()
@@ -581,12 +576,7 @@ def _setup_reference_model(
 
         # disabling dropout if found - non-determinism leads to issues in e.g. comparing logprobs
         # between ref policy and current policy
-        for module in model.modules():
-            if isinstance(module, torch.nn.Dropout):
-                warn(
-                    f"Dropout found in {module}. This is likely to cause issues during training. Disabling."
-                )
-                module.p = 0
+        disable_dropout(model)
 
         for p in model.parameters():
             p.requires_grad = False

recipes/ppo_full_finetune_single_device.py

@@ -24,7 +24,7 @@
 from torchtune.modules import local_kv_cache
 from torchtune.recipe_interfaces import FTRecipeInterface
 from torchtune.rlhf import PPOStats, Trajectory
-from torchtune.training import DummyProfiler, PROFILER_KEY
+from torchtune.training import disable_dropout, DummyProfiler, PROFILER_KEY
 from tqdm import tqdm
 
 log = utils.get_logger("DEBUG")
@@ -568,20 +568,10 @@ def _setup_models(
 
         # disabling dropout if found - non-determinism leads to issues in e.g. comparing logprobs
         # between ref policy and current policy
-        for module in policy_model.modules():
-            if isinstance(module, torch.nn.Dropout):
-                warn(
-                    f"Dropout found in {module}. This is likely to cause issues during training. Disabling."
-                )
-                module.p = 0
-        for module in value_model.modules():
-            if isinstance(module, torch.nn.Dropout):
-                warn(
-                    f"Dropout found in {module}. This is likely to cause issues during training. Disabling."
-                )
-                module.p = 0
+        disable_dropout(policy_model)
+        disable_dropout(value_model)
 
-        # disabling grad and dropout in reward and reference policy models
+        # disabling grad in reward and reference policy models
         reward_model.eval()
         ref_policy_model.eval()
 

tests/torchtune/training/test_model_util.py

@@ -7,7 +7,7 @@
 import warnings
 
 import torch
-from torchtune.training.model_util import disable_dropout
+from torchtune.training._model_util import disable_dropout
 
 
 class TestDisableDropout:
@@ -28,17 +28,32 @@ def test_disable_dropout(self):
 
     def test_disable_dropout_warning(self):
         """
-        Tests that a warning is issued when dropout layers are found and disabled.
+        Tests that correct number warning is issued when dropout layers are found and disabled.
         """
         model = torch.nn.Sequential(
             torch.nn.Linear(10, 10),
             torch.nn.Dropout(p=0.5),
             torch.nn.ReLU(),
             torch.nn.Dropout(p=0.3),
+            torch.nn.Dropout(p=0.0),
         )
         with warnings.catch_warnings(record=True) as w:
             warnings.simplefilter("always")
             disable_dropout(model)
             assert len(w) == 2, "Expected 2 warnings for 2 dropout layers."
             assert issubclass(w[-1].category, UserWarning)
-            assert "Dropout found in" in str(w[-1].message)
+            assert "Found Dropout with" in str(w[-1].message)
+
+    def test_disable_dropout_no_warning(self):
+        """
+        Tests that no warning is issued when there are no dropout layers.
+        """
+        model = torch.nn.Sequential(
+            torch.nn.Linear(10, 10),
+            torch.nn.ReLU(),
+            torch.nn.Linear(10, 10),
+        )
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter("always")
+            disable_dropout(model)
+            assert len(w) == 0, "Expected no warnings when there are no dropout layers."

torchtune/training/__init__.py

@@ -25,6 +25,7 @@
     validate_no_params_on_meta_device,
 )
 from torchtune.training._grad_scaler import scale_grads
+from torchtune.training._model_util import disable_dropout
 from torchtune.training._profiler import (
     DEFAULT_PROFILE_DIR,
     DEFAULT_PROFILER_ACTIVITIES,
@@ -135,4 +136,5 @@
     "OffloadActivations",
     "FormattedCheckpointFiles",
     "scale_grads",
+    "disable_dropout",
 ]

torchtune/training/model_util.py torchtune/training/_model_util.pytorchtune/training/model_util.py renamed to torchtune/training/_model_util.py

@@ -16,8 +16,8 @@ def disable_dropout(model: torch.nn.Module) -> None:
         model (torch.nn.Module): The model in which dropout layers should be disabled.
     """
     for module in model.modules():
-        if isinstance(module, torch.nn.Dropout):
+        if isinstance(module, torch.nn.Dropout) and module.p != 0:
             warnings.warn(
-                f"Dropout found in {module}. This is likely to cause issues during training. Disabling."
+                f"Found Dropout with value {module.p} in module {module}. Setting to zero."
             )
             module.p = 0