From 164db756de656153d574bd6712b8c405be9ec228 Mon Sep 17 00:00:00 2001
From: Daniel King <daniel@mosaicml.com>
Date: Wed, 15 Feb 2023 11:43:39 -0800
Subject: [PATCH 1/2] wip fsdp test

---
 tests/models/test_hf_model.py | 37 ++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/tests/models/test_hf_model.py b/tests/models/test_hf_model.py
index 43fb782a3a..f5060cc766 100644
--- a/tests/models/test_hf_model.py
+++ b/tests/models/test_hf_model.py
@@ -5,19 +5,20 @@
 import os
 import tempfile
 from pathlib import Path
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 from unittest.mock import patch
 from urllib.parse import urlparse
 
 import pytest
 import torch
+from packaging import version
 from torch.utils.data import DataLoader
 from torchmetrics import Metric
 from torchmetrics.classification import Accuracy
 
 from composer.metrics.nlp import LanguageCrossEntropy, MaskedAccuracy
 from composer.trainer import Trainer
-from composer.utils import dist
+from composer.utils import dist, is_model_fsdp
 from tests.common.datasets import RandomTextClassificationDataset, RandomTextLMDataset
 from tests.common.models import (configure_tiny_bert_model, configure_tiny_bert_tokenizer, configure_tiny_gpt2_model,
                                  configure_tiny_gpt2_tokenizer)
@@ -231,7 +232,8 @@ def get_lm_trainer(hf_model,
                    save_folder,
                    load_path: Optional[str] = None,
                    is_conditional_generation: bool = False,
-                   do_eval: bool = False):
+                   do_eval: bool = False,
+                   fsdp_config: Optional[Dict[str, Any]] = None):
     transformers = pytest.importorskip('transformers')
     from composer.models import HuggingFaceModel
 
@@ -280,7 +282,8 @@ def get_lm_trainer(hf_model,
                       save_folder=save_folder,
                       save_interval='1ep',
                       save_filename='hf-checkpoint.pt',
-                      load_path=load_path)
+                      load_path=load_path,
+                      fsdp_config=fsdp_config)
     return trainer
 
 
@@ -540,11 +543,31 @@ def test_encoder_decoder(tiny_t5_model, tiny_t5_tokenizer):
     trainer.eval()
 
 
-def test_hf_return_dict_false(tiny_bert_config, tiny_bert_tokenizer):
+@pytest.mark.gpu
+@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.13.0'),
+                    reason='requires PyTorch 1.13 or higher')
+def test_hf_fsdp(tiny_bert_config, tiny_bert_tokenizer):
     transformers = pytest.importorskip('transformers')
 
-    tiny_bert_config.return_dict = False
     tiny_bert_model = transformers.AutoModelForMaskedLM.from_config(tiny_bert_config)
-    trainer = get_lm_trainer(tiny_bert_model, tiny_bert_tokenizer, None, do_eval=True)
 
+    fsdp_config = {
+        'sharding_strategy': 'FULL_SHARD',
+        'min_params': 1e8,
+        'cpu_offload': False,
+        'mixed_precision': 'PURE',
+        'backward_prefetch': 'BACKWARD_PRE',
+        'activation_checkpointing': False,
+        'activation_cpu_offload': False,
+        'verbose': False
+    }
+
+    trainer = get_lm_trainer(tiny_bert_model, tiny_bert_tokenizer, None, fsdp_config=fsdp_config)
+
+    print(trainer.state.model)
+    assert False
+
+    assert is_model_fsdp(trainer.state.model)
+
+    assert trainer.state.fsdp_enabled
     trainer.fit()

From a72b968c3e1b343c4f37430ad4e41f0af87f5c3a Mon Sep 17 00:00:00 2001
From: Daniel King <daniel@mosaicml.com>
Date: Wed, 15 Feb 2023 19:58:28 +0000
Subject: [PATCH 2/2] add test for hf + fsdp

---
 tests/models/test_hf_model.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/models/test_hf_model.py b/tests/models/test_hf_model.py
index f5060cc766..e50e02312c 100644
--- a/tests/models/test_hf_model.py
+++ b/tests/models/test_hf_model.py
@@ -546,6 +546,7 @@ def test_encoder_decoder(tiny_t5_model, tiny_t5_tokenizer):
 @pytest.mark.gpu
 @pytest.mark.skipif(version.parse(torch.__version__) < version.parse('1.13.0'),
                     reason='requires PyTorch 1.13 or higher')
+@pytest.mark.filterwarnings('ignore::UserWarning')
 def test_hf_fsdp(tiny_bert_config, tiny_bert_tokenizer):
     transformers = pytest.importorskip('transformers')
 
@@ -564,9 +565,6 @@ def test_hf_fsdp(tiny_bert_config, tiny_bert_tokenizer):
 
     trainer = get_lm_trainer(tiny_bert_model, tiny_bert_tokenizer, None, fsdp_config=fsdp_config)
 
-    print(trainer.state.model)
-    assert False
-
     assert is_model_fsdp(trainer.state.model)
 
     assert trainer.state.fsdp_enabled