Remove flatten params

CODEOWNERS

@@ -17,11 +17,11 @@
 # as an owner for all sections, so anyone on Composer Eng can approve any Composer PR
 # According to the CODEOWNER docs, the last match takes precedence, so @mosaicml/composer-team-eng
 # must be mentioned for each rule below.
-/composer/algorithms/ @dskhudia @mvpatel2000 @nik-mosaic
+/composer/algorithms/ @mosaicml/composer-team-eng
 /composer/cli/ @mosaicml/composer-team-eng
 /composer/datasets/ @mosaicml/composer-team-eng
-/composer/functional/ @dblalock @mvpatel2000
-/composer/loggers/ @eracah @dakinggg
+/composer/functional/ @mosaicml/composer-team-eng @dblalock
+/composer/loggers/ @mosaicml/composer-team-eng @eracah @dakinggg
 /composer/loss/ @mosaicml/composer-team-eng
 /composer/metrics/ @mosaicml/composer-team-eng
 /composer/models/ @mosaicml/composer-team-eng

composer/algorithms/stochastic_depth/README.md

@@ -16,7 +16,7 @@
 from torch.utils.data import DataLoader
 from tests.common import RandomImageDataset
 
-train_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
+train_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
 ```
 -->
 <!--pytest-codeblocks:cont-->
@@ -65,8 +65,8 @@ for epoch in range(1):
 from torch.utils.data import DataLoader
 from tests.common import RandomImageDataset
 
-train_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
-eval_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
+train_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
+eval_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
 ```
 -->
 <!--pytest-codeblocks:cont-->

composer/algorithms/weight_standardization/README.md

@@ -20,7 +20,7 @@ from torchvision import models
 from torch.utils.data import DataLoader
 from tests.common import RandomImageDataset
 
-my_train_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
+my_train_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
 my_cnn_model = models.resnet18()
 ```
 -->
@@ -63,8 +63,8 @@ from torch.utils.data import DataLoader
 from tests.common import RandomImageDataset
 
 cnn_composer_model = composer_resnet('resnet50')
-my_train_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
-my_eval_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
+my_train_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
+my_eval_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
 ```
 -->
 <!--pytest-codeblocks:cont-->

composer/trainer/dist_strategy.py

@@ -135,7 +135,6 @@ def set_fsdp_default(fsdp_config: Dict[str, Any]):
     fsdp_config.setdefault('backward_prefetch', 'BACKWARD_POST')
     fsdp_config.setdefault('backward_prefetch_limit', 1)
     fsdp_config.setdefault('cpu_offload', False)
-    fsdp_config.setdefault('flatten_parameters', True)
     fsdp_config.setdefault('forward_prefetch', False)
     fsdp_config.setdefault('forward_prefetch_limit', 1)
     fsdp_config.setdefault('ignored_modules', None)

docs/source/method_cards/stochastic_depth.md

@@ -16,7 +16,7 @@ Block-wise stochastic depth assigns every residual block a probability of droppi
 from torch.utils.data import DataLoader
 from tests.common import RandomImageDataset
 
-train_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
+train_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
 ```
 -->
 <!--pytest-codeblocks:cont-->
@@ -65,8 +65,8 @@ for epoch in range(1):
 from torch.utils.data import DataLoader
 from tests.common import RandomImageDataset
 
-train_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
-eval_dataloader = DataLoader(RandomImageDataset(), batch_size=2)
+train_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
+eval_dataloader = DataLoader(RandomImageDataset(size=2), batch_size=2)
 ```
 -->
 <!--pytest-codeblocks:cont-->

docs/source/notes/distributed_training.rst

@@ -188,7 +188,6 @@ The full spec and defaults for Composer's `fsdp_config` is here:
       'activation_cpu_offload': bool = True | False, # Default: False
       'backward_prefetch': str = 'BACKWARD_PRE' | 'BACKWARD_POST' | 'NONE', # Default: 'BACKWARD_POST'
       'cpu_offload': bool = True | False, # Default: False, cpu_offload not supported yet
-      'flatten_parameters': bool = True | False, # Default: True
       'forward_prefetch': bool = True | False, # Default: False
       'ignored_modules': Optional[Iterable[torch.nn.Module]], # Default: None
       'keep_low_precision_grads': bool = True | False, # Default: False

tests/conftest.py

@@ -111,14 +111,18 @@ def pytest_configure():
     if TRANSFORMERS_INSTALLED:
         from tests.fixtures.fixtures import (tiny_bert_config_helper, tiny_bert_model_helper,
                                              tiny_bert_tokenizer_helper, tiny_gpt2_config_helper,
-                                             tiny_gpt2_model_helper, tiny_gpt2_tokenizer_helper, tiny_t5_config_helper,
+                                             tiny_gpt2_model_helper, tiny_gpt2_tokenizer_helper, tiny_opt_config_helper,
+                                             tiny_opt_model_helper, tiny_opt_tokenizer_helper, tiny_t5_config_helper,
                                              tiny_t5_model_helper, tiny_t5_tokenizer_helper)
         pytest.tiny_bert_config = tiny_bert_config_helper()  # type: ignore
         pytest.tiny_bert_model = tiny_bert_model_helper(pytest.tiny_bert_config)  # type: ignore
         pytest.tiny_bert_tokenizer = tiny_bert_tokenizer_helper()  # type: ignore
         pytest.tiny_gpt2_config = tiny_gpt2_config_helper()  # type: ignore
         pytest.tiny_gpt2_model = tiny_gpt2_model_helper(pytest.tiny_gpt2_config)  # type: ignore
         pytest.tiny_gpt2_tokenizer = tiny_gpt2_tokenizer_helper()  # type: ignore
+        pytest.tiny_opt_config = tiny_opt_config_helper()  # type: ignore
+        pytest.tiny_opt_model = tiny_opt_model_helper(pytest.tiny_opt_config)  # type: ignore
+        pytest.tiny_opt_tokenizer = tiny_opt_tokenizer_helper()  # type: ignore
         pytest.tiny_t5_config = tiny_t5_config_helper()  # type: ignore
         pytest.tiny_t5_model = tiny_t5_model_helper(pytest.tiny_t5_config)  # type: ignore
         pytest.tiny_t5_tokenizer = tiny_t5_tokenizer_helper()  # type: ignore