refactor memroy snapshot (mosaicml#2960)

composer/callbacks/memory_snapshot.py

@@ -95,12 +95,12 @@ def __init__(
         else:
             self.remote_path_in_bucket = None
 
-        if version.parse(torch.__version__) > version.parse('2.1.0.dev'):  # type: ignore
-            # memory snapshot  is only supported in torch v2.1.0-rc1 or higher
+        if version.parse(torch.__version__.split('.dev')[0]) >= version.parse('2.1.0'):  # type: ignore
+            # MemorySnapshot is only supported in torch v2.1.0-rc1 or higher
             self._enabled = True
         else:
             self._enabled = False
-            log.warning('Memory snapshot is supported after PyTorch 2.1.0. Skipping memory snapshot callback.')
+            warnings.warn('Memory snapshot is supported after PyTorch 2.1.0. Skipping memory snapshot callback.')
 
     def init(self, state: State, logger: Logger) -> None:
         if not self._enabled:

tests/callbacks/test_memory_snapshot.py

@@ -12,17 +12,13 @@
 from composer.callbacks import MemorySnapshot
 from composer.loggers import LoggerDestination
 from composer.trainer import Trainer
-from tests.common import RandomClassificationDataset, SimpleModel, device
+from tests.common import RandomClassificationDataset, SimpleModel
 
 
-@device('cpu', 'gpu')
-def test_memory_snapshot_warnings_on_cpu_models(device: str):
-    if version.parse(torch.__version__) <= version.parse('2.1.0.dev'):
-        # memory snapshot is supported after PyTorch 2.1.0.
-        return
-    # Error if the user sets device=cpu even when cuda is available
-    del device  # unused. always using cpu
-    with pytest.warns(UserWarning, match='The memory snapshot only works on CUDA devices'):
+@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.1.0'),
+                    reason='OOM Observer requires PyTorch 2.1 or higher')
+def test_memory_snapshot_warnings_on_cpu_models():
+    with pytest.warns(UserWarning):
         Trainer(
             model=SimpleModel(),
             callbacks=MemorySnapshot(),
@@ -44,16 +40,13 @@ def upload_file(self, state: State, remote_file_name: str, file_path: pathlib.Pa
 
 @pytest.mark.gpu
 @pytest.mark.parametrize('interval', ['1ba'])
+@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.1.0'),
+                    reason='OOM Observer requires PyTorch 2.1 or higher')
 def test_memory_snapshot(interval: str):
-    if version.parse(torch.__version__) <= version.parse('2.1.0.dev'):
-        # memory snapshot is supported after PyTorch 2.1.0.
-        return
     # Construct the callbacks
     skip_batches = 0
     memory_snapshot = MemorySnapshot(skip_batches=skip_batches, interval=interval)
-
     simple_model = SimpleModel()
-
     file_tracker_destination = FileUploaderTracker()
 
     # Construct the trainer and train