Repository: mosaicml/composer

Issue:
Single GPU models don't load weights from checkpoint
** To reproduce

Steps to reproduce the behavior:
```
import torch
from composer.trainer import Trainer
from composer import models

model = models.CIFAR10_ResNet56()
trainer = Trainer(model=model, train_dataloader=train_dataloader, eval_dataloader=None, max_duration="1ep", save_folder="checkpoints")
trainer.checkpoint_saver.save_checkpoint(trainer.state, trainer.seed, trainer.device)
new_model = models.CIFAR10_ResNet56()
new_trainer = Trainer(model=new_model, train_dataloader=train_dataloader, eval_dataloader=None, max_duration="1ep", load_path=os.path.join(trainer.checkpoint_saver.checkpoint_folder, "ep1.tar"))
assert torch.allclose(model.module.conv.weight, new_model.module.conv.weight)
```

## Expected behavior

The new model should have received the checkpoint's weights during the init of `new_trainer`. It doesn't and the reason seems to be due to here: https://github.com/mosaicml/composer/blob/dev/composer/core/state.py#L296-L302. If DDP is not enabled, the `state_dict` is never restored for the model.




Discussion:

Resolving patch:
diff --git a/composer/core/state.py b/composer/core/state.py
--- a/composer/core/state.py
+++ b/composer/core/state.py
@@ -295,11 +295,11 @@ def load_model_state(self, state_dict: types.StateDict, strict: bool):
         """
         if state_dict["_is_model_ddp_wrapped"] and not isinstance(self.model, DistributedDataParallel):
             torch.nn.modules.utils.consume_prefix_in_state_dict_if_present(state_dict['model'], "module.")
-            missing_keys, unexpected_keys = self.model.load_state_dict(state_dict['model'], strict=strict)
-            if len(missing_keys) > 0:
-                logger.warning(f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}")
-            if len(unexpected_keys) > 0:
-                logger.warning(f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}")
+        missing_keys, unexpected_keys = self.model.load_state_dict(state_dict['model'], strict=strict)
+        if len(missing_keys) > 0:
+            logger.warning(f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}")
+        if len(unexpected_keys) > 0:
+            logger.warning(f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}")
 
     def load_state_dict(self, state: types.StateDict, strict: bool = False):
         """Loads the state.
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -232,15 +232,14 @@ def __init__(
 
         self.find_unused_parameters = find_unused_parameters
 
-        if self.deepspeed_enabled:
-            import deepspeed
-            deepspeed.init_distributed()
-        else:
+        if self.deepspeed_enabled or dist.get_world_size() > 1:
+            # deepspeed requires torch.distributed to be initialized, even if the world size is 1
+            # distributed is always required with multi-rank training
             dist.initialize_dist(self.device.dist_backend, datetime.timedelta(seconds=dist_timeout))
-            if ddp_sync_strategy is None:
-                self.ddp_sync_strategy = DDPSyncStrategy.SINGLE_AUTO_SYNC if not find_unused_parameters else DDPSyncStrategy.FORCED_SYNC
-            else:
-                self.ddp_sync_strategy = DDPSyncStrategy(ddp_sync_strategy)
+        if ddp_sync_strategy is None:
+            self.ddp_sync_strategy = DDPSyncStrategy.SINGLE_AUTO_SYNC if not find_unused_parameters else DDPSyncStrategy.FORCED_SYNC
+        else:
+            self.ddp_sync_strategy = DDPSyncStrategy(ddp_sync_strategy)
 
         # `eval_dataloader` could be a dataloader, dataspec, evaluator, List[Evaluator], Tuple[Evaluator, ...], or dict of Dataspec hparams
         # convert it to `List[Evaluator]`
diff --git a/composer/utils/dist.py b/composer/utils/dist.py
--- a/composer/utils/dist.py
+++ b/composer/utils/dist.py
@@ -4,6 +4,7 @@
 
 import datetime
 import os
+import textwrap
 import warnings
 from typing import Any, List, Optional, Sequence, TypeVar, cast
 
@@ -14,10 +15,12 @@
 TObj = TypeVar("TObj")
 
 
-def _get_distributed_config_var(env_var: str,
-                                human_name: str,
-                                default: int,
-                                fetch_fn_name: Optional[str] = None) -> int:
+def _get_distributed_config_var(
+    env_var: str,
+    human_name: str,
+    default: int,
+    fetch_fn_name: Optional[str] = None,
+) -> int:
     if not dist.is_available():
         warnings.warn("DistributedDefaultValueWarning: Torch distributed is not available; "
                       f"returning {default} for {human_name}")
@@ -72,7 +75,7 @@ def get_local_world_size() -> int:
     Returns:
         int: The local world size
     """
-    return _get_distributed_config_var(env_var="LOCAL_WORLD_SIZE", human_name="local world size", default=1)
+    return _get_distributed_config_var(env_var="LOCAL_WORLD_SIZE", default=1, human_name="local world size")
 
 
 def get_local_rank() -> int:
@@ -81,7 +84,7 @@ def get_local_rank() -> int:
     Returns:
         int: The local world size
     """
-    return _get_distributed_config_var(env_var="LOCAL_RANK", human_name="local rank", default=0)
+    return _get_distributed_config_var(env_var="LOCAL_RANK", default=0, human_name="local rank")
 
 
 def get_node_rank() -> int:
@@ -91,7 +94,7 @@ def get_node_rank() -> int:
     Returns:
         int: The node rank, starting at 0.
     """
-    return _get_distributed_config_var(env_var="NODE_RANK", human_name="node rank", default=0)
+    return _get_distributed_config_var(env_var="NODE_RANK", default=0, human_name="node rank")
 
 
 def barrier() -> None:
@@ -241,12 +244,27 @@ def initialize_dist(backend: str, timeout: datetime.timedelta):
                                "wish to change backends, please restart the python process.")
         return
 
-    if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
-        warnings.warn("NoDistributedWarning: RANK and WORLD_SIZE env vars not set; assuming no "
-                      "parallelization. If this is unexpected, make sure you are running your "
-                      "training script with the composer CLI tool.")
-    elif get_world_size() > 1:
-        dist.init_process_group(backend, timeout=timeout)
+    dist_env_variable_names = ("NODE_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "RANK", "LOCAL_RANK")
+
+    is_missing_all_dist_env_vars = all(x not in os.environ for x in dist_env_variable_names)
+    if is_missing_all_dist_env_vars:
+        # missing all variables, in which case we should assume a single process
+        # if any variables are set, then it's likely an incomplete configuration, in which case we should not assume
+        # defaults (it would be better to let dist.init_process_group crash)
+        warnings.warn(
+            textwrap.dedent(f"""\
+                NoDistributedWarning: No distributed environment variables are set; assuming no
+                parallelization. If this is unexpected, please run the script with the composer CLI tool."""))
+        # setting the environment variables to single-rank defaults
+        os.environ["LOCAL_RANK"] = "0"
+        os.environ["RANK"] = "0"
+        os.environ["LOCAL_WORLD_SIZE"] = "1"
+        os.environ["WORLD_SIZE"] = "1"
+        os.environ["NODE_RANK"] = "0"
+        dist.init_process_group(backend, store=dist.HashStore(), world_size=1, rank=0)
+        return
+
+    dist.init_process_group(backend, timeout=timeout)
 
 
 def get_sampler(dataset, *, drop_last: bool, shuffle: bool) -> torch.utils.data.Sampler:


Test patch:
diff --git a/scripts/test.sh b/scripts/test.sh
--- a/scripts/test.sh
+++ b/scripts/test.sh
@@ -11,5 +11,6 @@ JUNIT_PREFIX=${JUNIT_PREFIX:-'build/output/composer'}
 
 mkdir -p $(dirname $JUNIT_PREFIX)
 
+python -m coverage run -m pytest --junitxml $JUNIT_PREFIX.n0.junit.xml $@  # identify non-ddp errors
 python -m composer.cli.launcher -n 1 --master_port 26000 -m coverage run -m pytest --junitxml $JUNIT_PREFIX.n1.junit.xml $@
 python -m composer.cli.launcher -n 2 --master_port 26000 -m coverage run -m pytest --junitxml $JUNIT_PREFIX.n2.junit.xml $@
diff --git a/tests/fixtures/distributed_fixtures.py b/tests/fixtures/distributed_fixtures.py
--- a/tests/fixtures/distributed_fixtures.py
+++ b/tests/fixtures/distributed_fixtures.py
@@ -31,26 +31,18 @@ def configure_dist(request: pytest.FixtureRequest):
     if is_deepspeed and is_gpu:
         pytest.fail('Tests should be marked as deepspeed or gpu, not both. Deepspeed tests will run on a gpu.')
 
-    if is_deepspeed:
-        assert not is_gpu
-        if not "RANK" in os.environ:
-            os.environ["RANK"] = str(0)
-            os.environ["LOCAL_RANK"] = str(0)
-            os.environ["WORLD_SIZE"] = str(1)
-            os.environ["MASTER_ADDR"] = "127.0.0.1"
-            os.environ["MASTER_PORT"] = str(26000)
-        import deepspeed
-        deepspeed.init_distributed(timeout=DIST_TIMEOUT)
+    dist_env_variable_names = ("NODE_RANK", "WORLD_SIZE", "LOCAL_WORLD_SIZE", "RANK", "LOCAL_RANK")
+
+    is_missing_all_dist_env_vars = all(x not in os.environ for x in dist_env_variable_names)
+
+    if is_missing_all_dist_env_vars and (not is_deepspeed):
+        # no distributed -- i.e. running pytest on the CLI
+        # deepspeed still requires distributed
         return
 
-    assert not is_deepspeed
-    backend = "nccl" if is_gpu else "gloo"
+    backend = "nccl" if (is_gpu or is_deepspeed) else "gloo"
     if not torch.distributed.is_initialized():
-        if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
-            torch.distributed.init_process_group(backend, timeout=DIST_TIMEOUT)
-        else:
-            store = torch.distributed.HashStore()
-            torch.distributed.init_process_group(backend, timeout=DIST_TIMEOUT, store=store, world_size=1, rank=0)
+        dist.initialize_dist(backend, timeout=DIST_TIMEOUT)
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/trainer/test_checkpoint.py b/tests/trainer/test_checkpoint.py
--- a/tests/trainer/test_checkpoint.py
+++ b/tests/trainer/test_checkpoint.py
@@ -96,7 +96,7 @@ def assert_weights_equivalent(original_trainer_hparams: TrainerHparams, new_trai
     recovered_weights = new_trainer.state.model.parameters()
 
     for p1, p2 in zip(original_weights, recovered_weights):
-        assert (p1.data.ne(p2.data).sum() == 0)
+        assert (p1.data == p2.data).all()
 
 
 @pytest.fixture
@@ -404,6 +404,7 @@ def validate_events_called_expected_number_of_times(trainer: Trainer):
 
 
 def test_checkpoint_load_uri(tmpdir: pathlib.Path):
+    pytest.xfail("example.com sometimes returns a 404. Need to mock out the actual download")
     loader = CheckpointLoader("https://example.com")
     loader._retrieve_checkpoint(destination_filepath=str(tmpdir / "example"), rank=0, ignore_not_found_errors=False)
     with open(str(tmpdir / "example"), "r") as f:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -8,6 +8,7 @@
 import torch
 import torch.distributed
 import yahp as hp
+from torch.nn.parallel import DistributedDataParallel
 
 from composer.algorithms import AlgorithmHparams
 from composer.algorithms.alibi.alibi import AlibiHparams
@@ -97,8 +98,10 @@ def test_trainer_determinism(composer_trainer_hparams: TrainerHparams, device: D
     composer_trainer_hparams.max_duration = "2ep"
 
     first_trainer = composer_trainer_hparams.initialize_object()
+    first_model = first_trainer.state.model
     first_trainer.fit()
-    first_model = first_trainer.state.model.module
+    if isinstance(first_model, DistributedDataParallel):
+        first_model = first_model.module
     assert isinstance(first_model, ComposerModel)
     assert first_trainer.state.train_dataloader is not None
     first_loss = get_total_loss(first_model, first_trainer.state.train_dataloader, first_trainer.device)
@@ -106,8 +109,10 @@ def test_trainer_determinism(composer_trainer_hparams: TrainerHparams, device: D
     # Second trainer must be created after fitting the first so that the
     # seeds get fully reset for the second training run
     second_trainer = composer_trainer_hparams.initialize_object()
+    second_model = second_trainer.state.model
     second_trainer.fit()
-    second_model = second_trainer.state.model.module
+    if isinstance(second_model, DistributedDataParallel):
+        second_model = second_model.module
     assert isinstance(second_model, ComposerModel)
     assert second_trainer.state.train_dataloader is not None
     second_loss = get_total_loss(second_model, second_trainer.state.train_dataloader, second_trainer.device)