diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 8b249a5ccf..a5e905ebc7 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -16,13 +16,13 @@ Example:
 -->
 
 # Before submitting
-- [ ] Have you read the [contributor guidelines](https://github.com/mosaicml/composer/blob/dev/CONTRIBUTING.md)?
+- [ ] Have you read the [contributor guidelines](https://github.com/mosaicml/composer/blob/main/CONTRIBUTING.md)?
 - [ ] Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
 - [ ] Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
 - [ ] Did you update any related docs and document your change?
-- [ ] Did you update any related tests and add any new tests related to your change? (see [testing](https://github.com/mosaicml/composer/blob/dev/CONTRIBUTING.md#running-tests))
+- [ ] Did you update any related tests and add any new tests related to your change? (see [testing](https://github.com/mosaicml/composer/blob/main/CONTRIBUTING.md#running-tests))
 - [ ] Did you run the tests locally to make sure they pass?
-- [ ] Did you run `pre-commit` on your change? (see the `pre-commit` section of [prerequisites](https://github.com/mosaicml/composer/blob/dev/CONTRIBUTING.md#prerequisites))
+- [ ] Did you run `pre-commit` on your change? (see the `pre-commit` section of [prerequisites](https://github.com/mosaicml/composer/blob/main/CONTRIBUTING.md#prerequisites))
 
 <!--
 Thanks so much for contributing to composer! We really appreciate it :)
diff --git a/.github/workflows/code-quality.yaml b/.github/workflows/code-quality.yaml
index 432e031cb4..5a2c86221f 100644
--- a/.github/workflows/code-quality.yaml
+++ b/.github/workflows/code-quality.yaml
@@ -2,7 +2,6 @@ name: Code Quality Checks
 on:
   push:
     branches:
-    - dev
     - main
     - release/**
   pull_request:
@@ -19,6 +18,7 @@ jobs:
   code-quality:
     runs-on: ubuntu-20.04
     timeout-minutes: 15
+    if: github.repository_owner == 'mosaicml'
     strategy:
       matrix:
         python_version:
diff --git a/.github/workflows/daily.yaml b/.github/workflows/daily.yaml
index aa97c755c8..b64e68d493 100644
--- a/.github/workflows/daily.yaml
+++ b/.github/workflows/daily.yaml
@@ -4,7 +4,6 @@ on:
   - cron: "30 2 * * *"  # 2:30 every day
   push:
     branches:
-    - dev
     - main
     - release/**
   workflow_dispatch:
@@ -77,13 +76,10 @@ jobs:
       pytest-command: ${{ matrix.pytest_command }}
       pytest-markers: ${{ matrix.markers }}
       composer_package_name: ${{ matrix.composer_package_name }}
-      pytest-wandb-entity: "mosaicml-public-integration-tests"
-      pytest-wandb-project: "integration-tests-${{ github.sha }}"
       safe_directory: composer
     secrets:
       aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
       aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
-      wandb-api-key: ${{ secrets.WANDB_API_KEY }}
       code-eval-device: ${{ secrets.CODE_EVAL_DEVICE }}
       code-eval-url: ${{ secrets.CODE_EVAL_URL }}
       code-eval-apikey: ${{ secrets.CODE_EVAL_APIKEY }}
diff --git a/.github/workflows/pr-docker.yaml b/.github/workflows/pr-docker.yaml
index 93f0b51be1..352eab881b 100644
--- a/.github/workflows/pr-docker.yaml
+++ b/.github/workflows/pr-docker.yaml
@@ -2,7 +2,6 @@ name: PR Docker
 on:
   pull_request:
     branches:
-    - dev
     - main
     - release/**
     paths:
diff --git a/.github/workflows/pr-gpu.yaml b/.github/workflows/pr-gpu.yaml
index f6de8908c1..392a2665c8 100644
--- a/.github/workflows/pr-gpu.yaml
+++ b/.github/workflows/pr-gpu.yaml
@@ -6,7 +6,7 @@ on:
 # or dev
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
-  cancel-in-progress: ${{ github.ref != 'refs/heads/main' && github.ref != 'refs/heads/dev' }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
 jobs:
   pytest-gpu-1:
     uses: mosaicml/ci-testing/.github/workflows/pytest-gpu.yaml@v0.0.9
diff --git a/.github/workflows/smoketest.yaml b/.github/workflows/smoketest.yaml
index b7bb09aaab..a0a6c445d2 100644
--- a/.github/workflows/smoketest.yaml
+++ b/.github/workflows/smoketest.yaml
@@ -2,7 +2,6 @@ name: Smoketest
 on:
   push:
     branches:
-    - dev
     - main
     - release/**
   pull_request:
@@ -20,6 +19,7 @@ jobs:
   smoketest:
     runs-on: ubuntu-20.04
     timeout-minutes: 10
+    if: github.repository_owner == 'mosaicml'
     strategy:
       matrix:
         python_version:
diff --git a/composer/_version.py b/composer/_version.py
index 82928466f9..a38b61a722 100644
--- a/composer/_version.py
+++ b/composer/_version.py
@@ -3,4 +3,4 @@
 
 """The Composer Version."""
 
-__version__ = '0.23.5'
+__version__ = '0.24.0.dev0'
diff --git a/composer/checkpoint/__init__.py b/composer/checkpoint/__init__.py
index 33162fc5e6..a79d0cea55 100644
--- a/composer/checkpoint/__init__.py
+++ b/composer/checkpoint/__init__.py
@@ -10,6 +10,7 @@
     get_optim_state_dict,
     get_resumption_state_dict,
 )
+from composer.checkpoint.upload import upload_file
 
 __all__ = [
     'get_model_state_dict',
@@ -17,4 +18,5 @@
     'get_metadata_state_dict',
     'get_resumption_state_dict',
     'download_monolithic_checkpoint',
+    'upload_file',
 ]
diff --git a/composer/checkpoint/download.py b/composer/checkpoint/download.py
index 01a80beb5f..2e786939c3 100644
--- a/composer/checkpoint/download.py
+++ b/composer/checkpoint/download.py
@@ -40,10 +40,11 @@ def download_file(
     if node_ranks is not None and node_rank not in node_ranks:
         return
 
+    if source_uri.endswith('.symlink'):
+        source_path = download_and_extract_symlink(source_uri)
+
     object_store = maybe_create_object_store_from_uri(source_uri)
     _, _, source_path = parse_uri(source_uri)
-    if source_uri.endswith('.symlink'):
-        source_path = extract_path_from_symlink(source_path, object_store)
     assert object_store is not None
 
     @retry(num_attempts=num_attempts)
@@ -58,6 +59,21 @@ def _download():
     log.debug(f'Finished downloading {source_path} to {destination_path}')
 
 
+def download_and_extract_symlink(source_uri) -> str:
+    """Downloads a symlink file from the specified URI and returns the path it points to.
+
+    Args:
+        source_uri (str): The URI to download the symlink from.
+
+    Returns:
+        str: The path the symlink points to.
+    """
+    object_store = maybe_create_object_store_from_uri(source_uri)
+    _, _, source_path = parse_uri(source_uri)
+    source_path = extract_path_from_symlink(source_path, object_store)
+    return source_path
+
+
 def download_monolithic_checkpoint(
     source_uri: str,
     destination_path: str,
diff --git a/composer/checkpoint/load.py b/composer/checkpoint/load.py
new file mode 100644
index 0000000000..3ad52fce3d
--- /dev/null
+++ b/composer/checkpoint/load.py
@@ -0,0 +1,670 @@
+# Copyright 2024 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""API for loading checkpoints."""
+
+import contextlib
+import logging
+import os
+import pickle
+import tempfile
+import textwrap
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.distributed.checkpoint as DCP
+from packaging import version
+from torch import nn
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.optim import Optimizer
+
+from composer.checkpoint.download import download_and_extract_symlink, download_monolithic_checkpoint
+from composer.checkpoint.state_dict import (
+    _cast_state_dict_to_precision,
+    _extract_keys_from_state_dict,
+    _is_model_fsdp,
+    _remove_keys_from_state_dict,
+    get_model_state_dict,
+    get_optim_state_dict,
+)
+from composer.core.state import State
+from composer.distributed import prepare_fsdp_module
+from composer.models import ComposerModel
+from composer.utils import FSDPConfig, dist, reproducibility
+from composer.utils.checkpoint import (
+    DistCPObjectStoreReader,
+    FileSystemReaderWithValidation,
+    _ensure_valid_checkpoint,
+    _torch_load_with_validation,
+)
+from composer.utils.file_helpers import is_uri
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class CheckpointLoadOptions:
+    """Options for loading a checkpoint.
+
+    Args:
+        load_model (bool): Whether to load the model checkpoint.
+        load_optimizer (bool): Whether to load the optimizer checkpoint.
+        load_resumption_state (bool): Whether to load the resumption state.
+        include_keys (Optional[Union[str, Sequence[str]]]): Specific keys to include in the state dict.
+        ignore_keys (Optional[Union[str, Sequence[str]]]): Specific keys to ignore in the state dict.
+        sharded_checkpoint (bool): Whether the checkpoint is sharded or not.
+        shard_as_needed_during_load (bool): Whether to shard the model and optimizer during load. If model is unsharded and checkpoint is sharded, model will be sharded before.
+            If model is unsharded and checkpoint is unsharded, model will be sharded after loading the checkpoint.
+        strict (bool): Whether to load the checkpoint strictly.
+        precision (str): The precision to cast the model/optim to after loading
+        cpu_offload (bool): Whether to offload the state dict to CPU before loading.
+        load_planner (Optional[LoadPlanner]): The load planner to use for loading the checkpoint.
+        fsdp_config (Optional[Union[FSDPConfig, dict]]): The FSDP config to use for sharding the model and optimizer.
+        device_mesh (Optional[DeviceMesh]): The device mesh to use for loading the checkpoint.
+        seed (Optional[int]): The seed to use for sharding the model and optimizer.
+    """
+    load_model: bool = True
+    load_optimizer: bool = False
+    load_resumption_state: bool = False  # dataset_state, timestamp,
+    # Specific key-level loading configs
+    # e.g.'model.layers.transformer.0.w1.weight'
+    include_keys: Optional[Union[str, Sequence[str]]] = None
+    ignore_keys: Optional[Union[str, Sequence[str]]] = None
+    sharded_checkpoint: bool = False  # TODO: Auto-detect sharded
+    shard_as_needed_during_load: bool = False
+    strict: bool = True
+    precision: str = 'fp32'
+    cpu_offload: bool = True
+    load_planner: Optional[Any] = None
+    fsdp_config: Optional[Union[FSDPConfig, dict]] = None
+    device_mesh: Optional[Any] = None
+    seed: Optional[int] = 42
+
+
+def load_checkpoint(
+    load_path: str,
+    load_options: Optional[Union[CheckpointLoadOptions, Dict]],
+    model: Optional[Union[ComposerModel, nn.Module]] = None,
+    optim: Optional[Optimizer] = None,
+    state: Optional[State] = None,
+    model_child_path: Optional[str] = None,
+    optim_child_path: Optional[str] = None,
+    resumption_filename: str = 'resumption.pkl',
+):
+    """Optionally download and load a checkpoint according to the options into specified state.
+
+    Args:
+        load_path (str): The path or uri to the checkpoint to load or symlink to the path/uri. If URI specified then files will be downloaded first.
+        load_options (CheckpointLoadOptions): The options for loading the checkpoint.
+        model (Optional[ComposerModel]): The model to load the checkpoint into. If not provided, the model from the state will be used.
+        optim (Optional[Optimizer]): The optimizer to load the checkpoint into. If not provided, the optimizer from the state will be used.
+        state (Optional[State]): The state to load the checkpoint. If not provided, the model and optimizer must be provided.
+        model_child_path (Optional[str]): The path to the model checkpoint within the load_path.
+            E.g. if your checkpoints are unsharded and saved in load_path/my_cool_model/my_model.pt, set model_child_path='my_cool_model/my_model.pt'. If not specified, 'model/model.pt' is used.
+            if your checkpoints are sharded and saved in load_path/my_sharded_model/ then set model_child_path='my_sharded_model'. If not specified, 'model' is used.
+        optim_child_path (Optional[str]): The path to the optimizer checkpoint within the load_path.
+            e.g. if your checkpoints are unsharded and saved in load_path/my_cool_optim/my_optim.pt, set optim_child_path='my_cool_optim/my_optim.pt'. If not specified, 'optim/optim.pt' is used.
+            if your checkpoints are sharded and saved in load_path/my_sharded_optim/ then set optim_child_path='my_sharded_optim'. If not specified, 'optim' is used.
+        resumption_filename (Optional[str]): The filename of the resumption state file. Default is 'resumption.pkl'.
+    """
+    if load_options is None:
+        load_options = CheckpointLoadOptions()
+    if isinstance(load_options, dict):
+        load_options = CheckpointLoadOptions(**load_options)
+
+    if load_options.load_model:
+        if model is None:
+            if state is None:
+                raise ValueError('Model or state must be provided if loading model checkpoint.')
+            model = state.model
+        if model_child_path is None:
+            model_child_path = 'model' if load_options.sharded_checkpoint else 'model/model.pt'
+
+    if load_options.load_optimizer:
+        if optim is None:
+            if state is None:
+                raise ValueError('Optimizer or state must be provided if loading optimizer checkpoint.')
+            optim = state.optimizers[0]
+        if optim_child_path is None:
+            optim_child_path = 'optim' if load_options.sharded_checkpoint else 'optim/optim.pt'
+
+    if load_options.load_resumption_state:
+        if state is None:
+            raise ValueError('State must be provided if loading resumption state.')
+        load_resumption_checkpoint(state, os.path.join(load_path, resumption_filename))
+
+    if load_options.load_model:
+        assert model is not None
+        assert model_child_path is not None
+        model_load_path = os.path.join(load_path, model_child_path)
+        if state is not None:
+            state.automicrobatch_fsdp_hook_handles, state.fsdp_modules = load_model_checkpoint(
+                model,
+                load_path=model_load_path,
+                load_options=load_options,
+            )
+        else:
+            load_model_checkpoint(
+                model,
+                load_path=model_load_path,
+                load_options=load_options,
+            )
+
+    if load_options.load_optimizer:
+        assert optim_child_path is not None
+        optim_load_path = os.path.join(load_path, optim_child_path)
+        assert model is not None
+        assert optim is not None
+        load_optim_checkpoint(
+            model,
+            optim,
+            optim_load_path,
+            load_options=load_options,
+        )
+
+
+def load_model_checkpoint(
+    model: Union[ComposerModel, nn.Module],
+    load_path: Optional[str] = None,
+    load_options: Optional[Union[CheckpointLoadOptions, Dict]] = None,
+    seed: int = 42,
+) -> Tuple[list, dict]:
+    """Load a a model checkpoint from the specified path into the model.
+
+    Args:
+        model (Union[ComposerModel, nn.Module]): The model to load the checkpoint into.
+        load_path (Optional[str]): The path to the checkpoint to load.
+        load_options (Optional[Union[CheckpointLoadOptions, Dict]]): The options for loading the checkpoint.
+        seed (int): The seed to use for sharding the model and optimizer.
+    """
+    if load_options is None:
+        load_options = CheckpointLoadOptions()
+    elif isinstance(load_options, dict):
+        load_options = CheckpointLoadOptions(**load_options)
+    if load_options.include_keys is not None and load_options.ignore_keys is not None:
+        raise ValueError('Only one of include_keys or ignore_keys can be specified.')
+
+    if load_options.include_keys is not None or load_options.ignore_keys is not None:
+        load_options.strict = False
+
+    automicrobatch_fsdp_hook_handles = []
+    fsdp_modules = {}
+
+    if load_options.sharded_checkpoint:
+        if not _is_model_fsdp(model):
+            if load_options.shard_as_needed_during_load:
+                automicrobatch_fsdp_hook_handles, fsdp_modules = _shard_with_fsdp(
+                    model,
+                    fsdp_config=load_options.fsdp_config,
+                    precision=load_options.precision,
+                    seed=seed,
+                )
+            else:
+                raise ValueError(
+                    'Model is not sharded but checkpoint is sharded. Please pass in a model wrapped with FSDP or set shard_model_as_needed_during_load to True.',
+                )
+        assert load_path is not None
+        _load_sharded_model_checkpoint(model, load_path=load_path, load_options=load_options)
+    else:
+        if dist.get_global_rank() == 0:
+            assert load_path is not None
+            _load_unsharded_model_checkpoint(model, load_path=load_path, load_options=load_options)
+
+        if load_options.shard_as_needed_during_load and not _is_model_fsdp(model):
+            if load_options.fsdp_config is None:
+                load_options.fsdp_config = {'sync_module_states': True}
+            elif isinstance(load_options.fsdp_config, dict):
+                load_options.fsdp_config.update({'sync_module_states': True})
+            else:
+                load_options.fsdp_config.sync_module_states = True
+            automicrobatch_fsdp_hook_handles, fsdp_modules = _shard_with_fsdp(
+                model,
+                fsdp_config=load_options.fsdp_config,
+                precision=load_options.precision,
+                seed=seed,
+            )
+    return automicrobatch_fsdp_hook_handles, fsdp_modules
+
+
+def _shard_with_fsdp(
+    model: Union[ComposerModel, nn.Module],
+    optimizer: Optional[Optimizer] = None,
+    fsdp_config: Optional[Union[FSDPConfig, dict]] = None,
+    precision: Optional[str] = None,
+    seed: int = 42,
+) -> Tuple[list, dict]:
+    if fsdp_config is None:
+        fsdp_config = FSDPConfig()
+    if isinstance(fsdp_config, dict):
+        fsdp_config = FSDPConfig(**fsdp_config)
+    with reproducibility.seed_context(seed):
+        automicrobatch_fsdp_hook_handles, fsdp_modules = prepare_fsdp_module(
+            model,
+            optimizers=optimizer,
+            fsdp_config=fsdp_config,
+            precision=precision,
+        )
+    return automicrobatch_fsdp_hook_handles, fsdp_modules
+
+
+def _load_sharded_model_checkpoint(
+    model: Union[ComposerModel, nn.Module],
+    load_path: str,
+    load_options: Optional[Union[CheckpointLoadOptions, Dict]],
+):
+    if load_options is None:
+        load_options = CheckpointLoadOptions()
+    elif isinstance(load_options, dict):
+        load_options = CheckpointLoadOptions(**load_options)
+    model_state_dict = get_model_state_dict(
+        model,
+        sharded_state_dict=True,
+        include_keys=load_options.include_keys,
+        ignore_keys=load_options.ignore_keys,
+    )
+    model_state_dict = download_and_load_sharded_state_dict(
+        load_path,
+        load_options.device_mesh,
+        model_state_dict,
+        load_options.load_planner,
+    )
+    model_state_dict = _cast_state_dict_to_precision(model_state_dict, load_options.precision)
+    # TODO: raise warning for unknown or missing keys.
+    log.debug(f'Loading sharded state dict into model.')
+    if version.parse(torch.__version__) >= version.parse('2.3.0') and dist.is_initialized():
+        torch_set_model_state_dict(
+            model,
+            model_state_dict,
+            strict=load_options.strict,
+            cpu_offload=load_options.cpu_offload,
+            sharded_state_dict=True,
+        )
+    else:
+        _load_model_state_dict_with_fsdp_context_manager(
+            model,
+            model_state_dict,
+            sharded_state_dict=True,
+            strict=load_options.strict,
+        )
+
+
+def _load_unsharded_model_checkpoint(
+    model: Union[ComposerModel, nn.Module],
+    load_path: str,
+    load_options: Optional[Union[CheckpointLoadOptions, Dict]] = None,
+):
+    if dist.get_global_rank() != 0:
+        return
+    if load_options is None:
+        load_options = CheckpointLoadOptions()
+    elif isinstance(load_options, dict):
+        load_options = CheckpointLoadOptions(**load_options)
+    if _is_model_fsdp(model):
+        raise ValueError(
+            'Model is sharded but checkpoint is not sharded. Please pass in a model not wrapped with FSDP.',
+        )
+    load_path_is_remote = is_uri(load_path)
+    download_dir_context = tempfile.TemporaryDirectory if load_path_is_remote else contextlib.nullcontext
+    with download_dir_context() as download_dir:
+        file_path = load_path
+        if load_path_is_remote:
+            filename = Path(load_path).name
+            assert download_dir is not None
+            file_path = os.path.join(download_dir, filename)
+            download_monolithic_checkpoint(load_path, download_dir)
+
+        model_state_dict = _torch_load_with_validation(file_path, map_location='cpu')
+        model_state_dict = _cast_state_dict_to_precision(model_state_dict, load_options.precision)
+        if load_options.include_keys is not None:
+            model_state_dict = _extract_keys_from_state_dict(model_state_dict, load_options.include_keys)
+        if load_options.ignore_keys is not None:
+            model_state_dict = _remove_keys_from_state_dict(model_state_dict, load_options.ignore_keys)
+        model.load_state_dict(model_state_dict, strict=load_options.strict)
+
+
+def load_optim_checkpoint(
+    model: Union[ComposerModel, nn.Module],
+    optim: torch.optim.Optimizer,
+    load_path: str,
+    load_options: Optional[Union[CheckpointLoadOptions, Dict]] = None,
+):
+    """Load an optimizer checkpoint from the specified path into the optimizer.
+
+    Args:
+        model (Union[ComposerModel, nn.Module]): The model to load the optimizer checkpoint into.
+        optim (torch.optim.Optimizer): The optimizer to load the checkpoint into.
+        load_path (Optional[str]): The path to the checkpoint to load.
+        load_options (Optional[Union[CheckpointLoadOptions, Dict]]): The options for loading the checkpoint.
+    """
+    if load_options is None:
+        load_options = CheckpointLoadOptions()
+    elif isinstance(load_options, dict):
+        load_options = CheckpointLoadOptions(**load_options)
+    if load_options.sharded_checkpoint and not _is_model_fsdp(model):
+        raise ValueError(
+            textwrap.dedent(
+                """Model/Optim is not sharded but checkpoint is sharded.
+                                         Please  pass in a model/optim wrapped with FSDP.""",
+            ),
+        )
+    if not load_options.sharded_checkpoint and not _is_model_fsdp(model) and load_options.shard_as_needed_during_load:
+        raise ValueError(
+            textwrap.dedent(
+                """Neither model nor optim nor checkpoint is sharded, but shard_as_needed_during_load is set.
+                                         Sharding the optim after load is not supported. please set shard_as_needed_during_load to False.
+                                         """,
+            ),
+        )
+
+    if load_options.shard_as_needed_during_load and _is_model_fsdp(model):
+        raise NotImplementedError(
+            'Loading optimizers for models that have been sharded during load (either before or after loading in the model checkpoint) is not currently supported.',
+        )
+
+    if load_options.sharded_checkpoint:
+        _load_sharded_optim_checkpoint(model=model, optim=optim, load_path=load_path, load_options=load_options)
+    else:
+        _load_unsharded_optim_checkpoint(
+            model=model,
+            optim=optim,
+            load_path=load_path,
+            precision=load_options.precision,
+        )
+
+
+def _load_sharded_optim_checkpoint(
+    model: Union[ComposerModel, nn.Module],
+    optim: torch.optim.Optimizer,
+    load_path: str,
+    load_options: CheckpointLoadOptions,
+):
+    if not _is_model_fsdp(model):
+        raise ValueError(
+            'Model is not sharded but checkpoint is sharded. Please either use load_model_checkpoint(model, load_path, optimizer, shard_as_needed_during_load=True) or pass in a model wrapped with FSDP.',
+        )
+    # if not _is_optimizer_sharded(optim):
+    #     raise ValueError("Optimizer is not sharded but checkpoint is sharded. Please pass in a sharded optimizer by passing a sharded model's parameters to an optimizer constructor.")
+    optim_state_dict = get_optim_state_dict(model, optim, sharded_state_dict=True)
+    optim_state_dict = download_and_load_sharded_state_dict(
+        load_path=load_path,
+        device_mesh=load_options.device_mesh,
+        state_dict=optim_state_dict,
+        load_planner=load_options.load_planner,
+    )
+    for param_key, param_state_dict in optim_state_dict['state'].items():
+        optim_state_dict['state'][param_key] = _cast_state_dict_to_precision(param_state_dict, load_options.precision)
+    if version.parse(torch.__version__) >= version.parse('2.3.0') and dist.is_initialized():
+        torch_set_optimizer_state_dict(
+            model,
+            optim,
+            optim_state_dict,
+            strict=load_options.strict,
+            cpu_offload=load_options.cpu_offload,
+            sharded_state_dict=True,
+        )
+    else:
+        _load_optim_state_dict_with_fsdp_context_manager(
+            model,
+            optim,
+            optim_state_dict,
+            sharded_state_dict=True,
+            strict=load_options.strict,
+        )
+
+
+def _load_unsharded_optim_checkpoint(
+    model: Union[ComposerModel, nn.Module],
+    optim: torch.optim.Optimizer,
+    load_path: str,
+    precision: str = 'fp32',
+):
+    if dist.get_global_rank() != 0:
+        return
+    if _is_model_fsdp(model):
+        raise ValueError('Model is sharded, but checkpoint is not sharded. Please pass in a model unwrapped from FSDP.')
+    # if _is_optimizer_sharded(optim):
+    #     raise ValueError("Optimizer is sharded, but checkpoint is not sharded. Please pass in an unsharded optimizer by passing an unsharded model's parameters to an optimizer constructor.")
+    load_path_is_remote = is_uri(load_path)
+    download_dir_context = tempfile.TemporaryDirectory if load_path_is_remote else contextlib.nullcontext
+    with download_dir_context() as download_dir:
+        file_path = load_path
+        if load_path_is_remote:
+            filename = Path(load_path).name
+            assert download_dir is not None
+            file_path = os.path.join(download_dir, filename)
+            download_monolithic_checkpoint(load_path, file_path)
+
+        optim_state_dict = _torch_load_with_validation(file_path, map_location='cpu')
+        for param_key, param_state_dict in optim_state_dict['state'].items():
+            optim_state_dict['state'][param_key] = _cast_state_dict_to_precision(param_state_dict, precision)
+        optim.load_state_dict(optim_state_dict)
+
+
+def _preprocess_local_load_path(load_path: str):
+    if os.path.islink(load_path):
+        load_path = os.path.join(os.path.dirname(load_path), os.readlink(load_path))
+    if os.path.exists(load_path):
+        if not os.path.isdir(load_path):
+            raise ValueError(f'load_path must be a directory when using sharded state dict. Got {load_path}')
+    else:
+        raise FileNotFoundError(f'{load_path} not found!')
+    return load_path
+
+
+def _load_model_state_dict_with_fsdp_context_manager(
+    model: nn.Module,
+    model_state_dict: dict,
+    sharded_state_dict: bool,
+    strict: bool,
+):
+    from torch.distributed.fsdp.fully_sharded_data_parallel import (
+        FullStateDictConfig,
+        ShardedStateDictConfig,
+        StateDictType,
+    )
+    state_dict_type = StateDictType.SHARDED_STATE_DICT if sharded_state_dict else StateDictType.FULL_STATE_DICT
+    state_dict_config = ShardedStateDictConfig(offload_to_cpu=True,) if sharded_state_dict else FullStateDictConfig(
+        offload_to_cpu=True,
+        rank0_only=True,
+    )
+    with FSDP.state_dict_type(model, state_dict_type=state_dict_type, state_dict_config=state_dict_config):
+        missing_keys, unexpected_keys = model.load_state_dict(
+            model_state_dict,
+            strict=strict,
+        )
+
+    if len(missing_keys) > 0:
+        log.warning(f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}")
+    if len(unexpected_keys) > 0:
+        log.warning(f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}")
+
+
+def _load_optim_state_dict_with_fsdp_context_manager(
+    model: nn.Module,
+    optim,
+    optim_state_dict: dict,
+    sharded_state_dict: bool,
+    strict: bool,
+):
+    from torch.distributed.fsdp.fully_sharded_data_parallel import (
+        FullOptimStateDictConfig,
+        ShardedOptimStateDictConfig,
+        StateDictType,
+    )
+    state_dict_type = StateDictType.SHARDED_STATE_DICT if sharded_state_dict else StateDictType.FULL_STATE_DICT
+    optim_state_dict_config = optim_state_dict_config = ShardedOptimStateDictConfig(
+    ) if sharded_state_dict else FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
+    with FSDP.state_dict_type(model, state_dict_type=state_dict_type, optim_state_dict_config=optim_state_dict_config):
+        optim_state_dict = FSDP.optim_state_dict_to_load(  #  type: ignore
+                                optim_state_dict=optim_state_dict, model=model, optim=optim,
+                            )
+    assert optim_state_dict is not None
+    optim.load_state_dict(optim_state_dict)
+
+
+def torch_set_model_state_dict(
+    model: torch.nn.Module,
+    model_state_dict: dict,
+    strict: bool,
+    cpu_offload: bool,
+    sharded_state_dict: bool = True,
+):
+    """Set the model state dict with the given options.
+
+    Args:
+        model (torch.nn.Module): The model to set the state dict to.
+        model_state_dict (dict): The model state dict to set.
+        strict (bool): Whether to set the state dict strictly.
+        cpu_offload (bool): Whether to offload the state dict to CPU before setting.
+        sharded_state_dict (bool): Whether the state dict is sharded or not.
+    """
+    if version.parse(torch.__version__) >= version.parse('2.3.0') and dist.is_initialized():
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
+        try:
+            set_model_state_dict(
+                model,
+                model_state_dict,
+                options=StateDictOptions(
+                    strict=strict,
+                    cpu_offload=cpu_offload,
+                    full_state_dict=not sharded_state_dict,
+                ),
+            )
+        except AttributeError as e:
+            # Issue: https://github.com/pytorch/pytorch/issues/127351
+            if "ShardedTensor' object has no attribute 'placements'" in str(e):
+                raise RuntimeError(
+                    textwrap.dedent(
+                        'PyTorch DTensor broke backwards compatibility in older checkpoints '
+                        'with ShardedTensor, which is now deprecated. To load old checkpoints, '
+                        'either downgrade to PyTorch <2.3.0 or explicitly pass process groups '
+                        'in the Trainer constructor via '
+                        "`parallelism_config = {'fsdp': {'process_group': 'mod1'}}`. We can "
+                        'provide assistance at https://github.com/mosaicml/composer/issues.',
+                    ),
+                ) from e
+            else:
+                raise e
+
+
+def torch_set_optimizer_state_dict(
+    model: torch.nn.Module,
+    optim: torch.optim.Optimizer,
+    optim_state_dict: dict,
+    strict: bool,
+    cpu_offload: bool,
+    sharded_state_dict: bool = True,
+):
+    """Set the optimizer state dict with the given options.
+
+    Args:
+        model (torch.nn.Module): The model to set the state dict to.
+        optim (torch.optim.Optimizer): The optimizer to set the state dict to.
+        optim_state_dict (dict): The optimizer state dict to set.
+        strict (bool): Whether to set the state dict strictly.
+        cpu_offload (bool): Whether to offload the state dict to CPU before setting.
+        sharded_state_dict (bool): Whether the state dict is sharded or not.
+    """
+    if version.parse(torch.__version__) >= version.parse('2.3.0') and dist.is_initialized():
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_optimizer_state_dict
+        set_optimizer_state_dict(
+            model=model,
+            optimizers=optim,
+            optim_state_dict=optim_state_dict,
+            options=StateDictOptions(
+                full_state_dict=not sharded_state_dict,
+                strict=strict,
+                cpu_offload=cpu_offload,
+            ),
+        )
+
+
+def download_and_load_sharded_state_dict(
+    load_path: str,
+    device_mesh: Optional[Any],
+    state_dict: dict,
+    load_planner: Optional[Any] = None,
+):
+    """A helper function to download and load a sharded state dict.
+
+    Args:
+        load_path (str): The path to the checkpoint to load.
+        device_mesh (Optional[Any]): The device mesh to use for loading the checkpoint.
+        state_dict (dict): The state dict to load.
+        load_planner (Optional[Any]): The load planner to use for loading the checkpoint.
+    """
+    load_path_is_remote = is_uri(load_path)
+    download_dir_context = tempfile.TemporaryDirectory if load_path_is_remote else contextlib.nullcontext
+    with download_dir_context() as download_dir:
+        if load_path_is_remote:
+            local_rank0_index = dist.get_global_rank() - dist.get_local_rank()
+            rank0_download_tempdir = str(dist.all_gather_object(download_dir)[local_rank0_index])
+            if load_path.endswith('.symlink'):
+                load_path = download_and_extract_symlink(load_path)
+            storage_reader = DistCPObjectStoreReader(
+                source_path=load_path,
+                destination_path=str(Path(rank0_download_tempdir) / Path('checkpoints')),
+                device_mesh=device_mesh,
+            )
+        else:
+            load_path = _preprocess_local_load_path(load_path)
+            storage_reader = FileSystemReaderWithValidation(load_path)
+
+            try:
+                storage_reader.read_metadata()
+            except AttributeError as e:
+                if '_MEM_FORMAT_ENCODING' in str(e):
+                    raise ValueError(
+                        'Unable to read checkpoint metadata. The checkpoint was likely saved with a '
+                        'newer version of torch. Upgrade your torch version to load this checkpoint.',
+                    )
+                else:
+                    raise
+
+        log.debug(f'Loading sharded state dict from {load_path} into memory.')
+        if version.parse(torch.__version__) < version.parse('2.2.0'):
+            DCP.load_state_dict(state_dict=state_dict, storage_reader=storage_reader, planner=load_planner)
+        else:
+            DCP.load(state_dict=state_dict, storage_reader=storage_reader, planner=load_planner)
+    return state_dict
+
+
+def load_resumption_checkpoint(state: State, load_path: str):
+    """Load the resumption state from the specified path into the state.
+
+    Args:
+        state (State): The state to load the resumption state into.
+        load_path (str): The path to the resumption state to load.
+
+    """
+    load_path = str(_ensure_valid_checkpoint(load_path))
+    resumption_state_dict = pickle.load(open(load_path, 'rb'))
+    if 'dataset_state' in resumption_state_dict:
+        state._load_dataset_state(resumption_state_dict['dataset_state'])
+    if 'timestamp' in resumption_state_dict:
+        state.timestamp.load_state_dict(resumption_state_dict['timestamp'])
+
+    if 'schedulers' in resumption_state_dict:
+        for scheduler in state.schedulers:
+            fqn = type(scheduler).__qualname__
+            scheduler.load_state_dict(resumption_state_dict['schedulers'][fqn])
+
+    if 'algorithms' in resumption_state_dict:
+        for algorithm in state.algorithms:
+            fqn = type(algorithm).__qualname__
+            algorithm.load_state_dict(dict(resumption_state_dict['algorithms'])[fqn])
+
+    if 'callbacks' in resumption_state_dict:
+        for callback in state.callbacks:
+            fqn = type(callback).__qualname__
+            callback.load_state_dict(dict(resumption_state_dict['callbacks'])[fqn])
+
+    if 'scaler' in resumption_state_dict:
+        fqn = type(state.scaler).__qualname__
+        assert state.scaler is not None
+        state.scaler.load_state_dict(resumption_state_dict['scaler'][fqn])
+
+    state.rank_zero_seed = resumption_state_dict['rank_zero_seed']
+    return resumption_state_dict['rng']
diff --git a/composer/checkpoint/save.py b/composer/checkpoint/save.py
index 03166d8802..edd2139004 100644
--- a/composer/checkpoint/save.py
+++ b/composer/checkpoint/save.py
@@ -65,7 +65,6 @@ class CheckpointSaveOptions:
         include_keys (Optional[Union[str, Sequence[str]]]): Keys to include in the saved model.
         ignore_keys (Optional[Union[str, Sequence[str]]]): Keys to ignore in the saved model.
     """
-    destination_dir: str
     save_frequency: Union[str, int, Time] = '1ep'
     dir_prefix: str = 'ep{epoch}-ba{batch}'
     overwrite: bool = False
@@ -75,15 +74,15 @@ class CheckpointSaveOptions:
     num_checkpoints_to_keep: int = -1
     save_format: str = 'pt'
     sharded_checkpoint: bool = False
-    precision: str = 'bf16'
+    precision: str = 'fp32'
     include_keys: Optional[Union[str, Sequence[str]]] = None
     ignore_keys: Optional[Union[str, Sequence[str]]] = None
 
 
 def save_checkpoint_to_disk(
+    destination_dir: str,
     state: State,
     options: Optional[Union[CheckpointSaveOptions, Dict]] = None,
-    destination_dir: Optional[str] = None,
 ):
     """Saves a checkpoint to disk.
 
@@ -92,24 +91,25 @@ def save_checkpoint_to_disk(
         options (Optional[Union[CheckpointSaveOptions, Dict]]): The options for saving the checkpoint.
             If None, destination_dir must be provided.
         destination_dir (Optional[str]): The directory to save the checkpoint to.
-            If options is provided, this will overwrite options.destination_dir.
+            If options is provided, this will overwrite save_path.
     """
     if options is None:
-        if destination_dir is None:
-            raise ValueError('destination_dir must be provided if options is None')
-        options = CheckpointSaveOptions(destination_dir=destination_dir)
+        options = CheckpointSaveOptions()
     else:
         if isinstance(options, Dict):
             options = CheckpointSaveOptions(**options)
-        if destination_dir is not None:
-            options.destination_dir = destination_dir
-    save_path = os.path.join(options.destination_dir, options.dir_prefix)
+
+    save_path = os.path.join(destination_dir, options.dir_prefix)
     save_path = format_name_with_dist_and_time(save_path, state.run_name, state.timestamp)
     os.makedirs(save_path, exist_ok=True)
     if options.save_model:
+        model_save_path = (
+            os.path.join(save_path, MODEL_CHECKPOINT_DIRECTORY_NAME) if options.sharded_checkpoint else
+            os.path.join(save_path, MODEL_CHECKPOINT_DIRECTORY_NAME, MONOLITHIC_MODEL_CHECKPOINT_FILENAME)
+        )
         save_model_to_disk(
             state.model,
-            save_path,
+            model_save_path,
             options.sharded_checkpoint,
             options.precision,
             options.include_keys,
@@ -118,27 +118,36 @@ def save_checkpoint_to_disk(
             options.save_format,
         )
     if options.save_optimizer:
+        optim_save_path = os.path.join(save_path,
+                                       OPTIM_CHECKPOINT_DIRECTORY_NAME) if options.sharded_checkpoint else os.path.join(
+                                           save_path,
+                                           OPTIM_CHECKPOINT_DIRECTORY_NAME,
+                                           OPTIM_MONO_CHECKPOINT_FILENAME,
+                                       )
         optimizer = state.optimizers[0]
         save_optim_to_disk(
             state.model,
             optimizer,
-            save_path,
+            optim_save_path,
             options.sharded_checkpoint,
             options.precision,
             options.overwrite,
             options.save_format,
         )
     if options.save_resumption_state:
-        save_resumption_state_to_disk(state, save_path)
+        resumption_save_path = os.path.join(save_path, RESUMPTION_CHECKPOINT_FILENAME)
+        save_resumption_state_to_disk(state, resumption_save_path)
 
+    md_save_path = os.path.join(save_path, METADATA_CHECKPOINT_FILENAME)
     save_composer_metadata_to_disk(
-        save_path,
+        md_save_path,
         state.model,
         options.sharded_checkpoint,
         options.precision,
         state.device,
         state.device_train_microbatch_size,
     )
+    return save_path
 
 
 def save_model_to_disk(
@@ -156,8 +165,8 @@ def save_model_to_disk(
     Args:
         model (Union[ComposerModel, torch.nn.Module]): The model to save.
         destination_dir (str): The directory to save the model to.
-            Model will be saved as distination_dir/models/model.pt if sharded_checkpoint is False,
-            otherwise all shards will be saved as destination_dir/models/__<rank>_0.distcp.
+            Model will be saved as distination_dir/model.pt if sharded_checkpoint is False,
+            otherwise all shards will be saved as destination_dir/__<rank>_0.distcp along with a metadata file (destination_dir/.metadata).
         sharded_checkpoint (bool): Whether to save the model as a sharded checkpoint.
         precision (str): The precision to save the model in. One of 'bf16', 'fp32', 'fp16', 'fp64'.
         include_keys (Optional[Union[str, Sequence[str]]]): Keys to include in the saved model.
@@ -179,14 +188,9 @@ def save_model_to_disk(
         include_keys,
         ignore_keys,
     )
-
-    destination_file_path = (
-        os.path.join(destination_dir, MODEL_CHECKPOINT_DIRECTORY_NAME) if sharded_checkpoint else
-        os.path.join(destination_dir, MODEL_CHECKPOINT_DIRECTORY_NAME, MONOLITHIC_MODEL_CHECKPOINT_FILENAME)
-    )
     saved_path = save_state_dict_to_disk(
         state_dict=model_state_dict,
-        destination_file_path=destination_file_path,
+        destination_file_path=destination_dir,
         overwrite=overwrite,
         save_format=save_format,
     )
@@ -208,8 +212,8 @@ def save_optim_to_disk(
         model (Union[ComposerModel, torch.nn.Module]): The model to save.
         optimizer (torch.optim.Optimizer): The optimizer to save.
         destination_dir (str): The directory to save the optimizer to.
-            Optimizer will be saved as destination_dir/optim/optim.pt if sharded_checkpoint is False,
-            otherwise all shards will be saved as destination_dir/optim/__<rank>_0.distcp.
+            Optimizer will be saved as destination_dir if sharded_checkpoint is False,
+            otherwise all shards will be saved as destination_dir/__<rank>_0.distcp along with a metadata file (destination_dir/.metadata).
         sharded_checkpoint (bool): Whether to save the optimizer as a sharded checkpoint.
         precision (str): The precision to save the optimizer in. One of 'bf16', 'fp32', 'fp16', 'fp64'.
         overwrite (bool): If True, the file will be overwritten if it exists.
@@ -221,15 +225,9 @@ def save_optim_to_disk(
         sharded_state_dict=sharded_checkpoint,
         precision=precision,
     )
-    destination_file_path = os.path.join(destination_dir,
-                                         OPTIM_CHECKPOINT_DIRECTORY_NAME) if sharded_checkpoint else os.path.join(
-                                             destination_dir,
-                                             OPTIM_CHECKPOINT_DIRECTORY_NAME,
-                                             OPTIM_MONO_CHECKPOINT_FILENAME,
-                                         )
     saved_path = save_state_dict_to_disk(
         state_dict=optim_state_dict,
-        destination_file_path=destination_file_path,
+        destination_file_path=destination_dir,
         overwrite=overwrite,
         save_format=save_format,
     )
@@ -238,7 +236,7 @@ def save_optim_to_disk(
 
 
 def save_composer_metadata_to_disk(
-    destination_dir: str,
+    destination_file_path: str,
     model: Optional[Union[ComposerModel, torch.nn.Module]] = None,
     sharded_state_dict: Optional[bool] = None,
     precision: Optional[Union[str, torch.dtype]] = None,
@@ -248,7 +246,7 @@ def save_composer_metadata_to_disk(
     """Saves metadata about the model to disk.
 
     Args:
-        destination_dir (str): The directory to save the metadata to.
+        destination_file_path (str): The path to save the metadata to.
         model (Optional[Union[ComposerModel, torch.nn.Module]]): The model to save metadata about.
         sharded_state_dict (Optional[bool]): Whether the model is sharded.
         precision (Optional[Union[str, torch.dtype]]): The precision of the model.
@@ -262,8 +260,7 @@ def save_composer_metadata_to_disk(
         device,
         device_train_microbatch_size,
     )
-    os.makedirs(destination_dir, exist_ok=True)
-    destination_file_path = os.path.join(destination_dir, METADATA_CHECKPOINT_FILENAME)
+    os.makedirs(str(Path(destination_file_path).parent), exist_ok=True)
 
     if dist.get_global_rank() == 0:
         with open(destination_file_path, 'w') as f:
@@ -273,16 +270,16 @@ def save_composer_metadata_to_disk(
 
 def save_resumption_state_to_disk(
     state: State,
-    destination_dir: str,
+    destination_file_path: str,
 ):
     """Saves the resumption state to disk.
 
     Args:
         state (State): The state to save.
-        destination_dir (str): The directory to save the resumption state to.
+        destination_file_path (str): The path to save the resumption state to.
     """
     resumption_state_dict = get_resumption_state_dict(state)
-    destination_file_path = os.path.join(destination_dir, RESUMPTION_CHECKPOINT_FILENAME)
+    os.makedirs(Path(destination_file_path).parent, exist_ok=True)
     with open(destination_file_path, 'wb') as f:
         pickle.dump(resumption_state_dict, f)
     return destination_file_path
diff --git a/composer/checkpoint/state_dict.py b/composer/checkpoint/state_dict.py
index 5f82836d7b..e36f31daf8 100644
--- a/composer/checkpoint/state_dict.py
+++ b/composer/checkpoint/state_dict.py
@@ -208,7 +208,7 @@ def get_resumption_state_dict(state: State) -> Dict[str, Any]:
     if state.scaler is not None:
         scaler_sd = _make_state_dict_for_list_of_objects(state.scaler)
         if scaler_sd != {}:
-            resumption_state_dict['scaler'] = state.scaler.state_dict()
+            resumption_state_dict['scaler'] = scaler_sd
 
     resumption_state_dict['rank_zero_seed'] = state.rank_zero_seed
     resumption_state_dict['run_name'] = state.run_name
diff --git a/composer/checkpoint/upload.py b/composer/checkpoint/upload.py
new file mode 100644
index 0000000000..343a7f8523
--- /dev/null
+++ b/composer/checkpoint/upload.py
@@ -0,0 +1,213 @@
+# Copyright 2024 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Useful functions for uploading checkpoints to remote object store."""
+
+import logging
+import os
+import pathlib
+import time
+from concurrent.futures import Future
+from typing import Optional
+
+import composer.utils.dist as dist
+from composer.core import Callback, State
+from composer.loggers import Logger, MosaicMLLogger
+from composer.utils import (
+    RemoteFilesExistingCheckStatus,
+    RemoteUploader,
+    create_symlink_file,
+    parse_uri,
+)
+
+log = logging.getLogger(__name__)
+
+
+class CheckpointUploadCallback(Callback):
+    """Callback class for async checkpoint uploading."""
+
+    def __init__(
+        self,
+        remote_uploader: RemoteUploader,
+        upload_timeout_in_seconds: int = 3600,
+    ):
+        self.remote_uploader = remote_uploader
+        if dist.get_global_rank() == 0:
+            self.rank_saves_symlinks = True
+        else:
+            self.rank_saves_symlinks = False
+
+        # Each Tuple contains information of:
+        # the future of checking if checkpoint files upload finish
+        # local symlink file name
+        # remote symlink file name
+        self.symlink_upload_tasks: list[tuple[Future[RemoteFilesExistingCheckStatus], str, Optional[str]]] = []
+        self.upload_timeout_in_seconds: int = upload_timeout_in_seconds
+
+        # Allow unit test to override this to make it faster
+        self._symlink_upload_wait_before_next_try_in_seconds = 30.0
+
+    def add_symlink_upload_task(
+        self,
+        remote_file_names: list[str],
+        symlink_file_name: Optional[str] = None,
+        symlink_remote_file_name: Optional[str] = None,
+    ):
+        check_remote_files_exist_future = None
+        if self.rank_saves_symlinks and symlink_file_name is not None:
+            check_remote_files_exist_future = self.remote_uploader.check_remote_files_exist_async(
+                remote_checkpoint_file_names=remote_file_names,
+                max_wait_time_in_seconds=self.upload_timeout_in_seconds,
+                wait_before_next_try_in_seconds=self._symlink_upload_wait_before_next_try_in_seconds,
+            )
+            self.symlink_upload_tasks.append(
+                (check_remote_files_exist_future, symlink_file_name, symlink_remote_file_name),
+            )
+
+    def _log_checkpoint_upload(self, logger: Logger):
+        for destination in logger.destinations:
+            if isinstance(destination, MosaicMLLogger):
+                destination.log_metadata({'checkpoint_uploaded_time': time.time()}, force_flush=True)
+
+    def batch_end(self, state: State, logger: Logger) -> None:
+        del state  # unused
+        self.remote_uploader.check_workers()
+        if not self.rank_saves_symlinks:
+            return
+        undone_symlink_upload_tasks = []
+        for (check_remote_files_exist_future, local_symlink_file,
+             remote_symlink_file) in reversed(self.symlink_upload_tasks):
+            if not check_remote_files_exist_future.done():
+                undone_symlink_upload_tasks.insert(
+                    0,
+                    (check_remote_files_exist_future, local_symlink_file, remote_symlink_file),
+                )
+                continue
+            else:
+                result = check_remote_files_exist_future.result()
+                if result == RemoteFilesExistingCheckStatus.EXIST:
+                    assert remote_symlink_file is not None
+                    self.remote_uploader.upload_file_async(
+                        remote_file_name=remote_symlink_file,
+                        file_path=pathlib.Path(local_symlink_file),
+                        overwrite=True,
+                    )
+                    self._log_checkpoint_upload(logger)
+                    break
+                else:
+                    raise RuntimeError(f'Failed to check if checkpoint files upload finish: {result}')
+        self.symlink_upload_tasks = undone_symlink_upload_tasks
+
+    def wait(self) -> None:
+        log.info('Waiting for checkpoint uploading to finish')
+        self.remote_uploader.wait()
+        if self.rank_saves_symlinks and len(self.symlink_upload_tasks) > 0:
+            log.debug('Uploading symlink to the latest checkpoint')
+            # We only need to upload a symlink pointing to the latest checkpoint files, so we can ignore successful uploads of older checkpoints.
+            check_remote_files_exist_future, local_symlink_file, remote_symlink_file = self.symlink_upload_tasks[-1]
+            result = check_remote_files_exist_future.result()
+            if result == RemoteFilesExistingCheckStatus.EXIST:
+                assert remote_symlink_file is not None
+                symlink_upload_future = self.remote_uploader.upload_file_async(
+                    remote_file_name=remote_symlink_file,
+                    file_path=pathlib.Path(local_symlink_file),
+                    overwrite=True,
+                )
+                symlink_upload_future.result()
+            else:
+                raise RuntimeError(f'Failed to check if checkpoint files upload finish: {result}')
+            self.symlink_upload_tasks = []
+        log.info('Checkpoint uploading finished!')
+
+    def fit_end(self, state: State, logger: Logger) -> None:
+        del state  # unused
+        self.wait()
+        self._log_checkpoint_upload(logger)
+
+    def post_close(self):
+        # Wait the symlink file upload to finish and close remote uploader
+        try:
+            self.remote_uploader.wait_and_close()
+        except Exception as e:
+            log.error(f'RemoteUploader run into exception {e}')
+
+
+
+def upload_file(
+    dest_dir: str,
+    source_path: Optional[str]=None,
+    symlink_granularity: Optional[str]=None, # file, dir, or None
+    symlink_name: Optional[str]='latest.symlink',
+    async_upload: bool = True,
+    state: Optional[State] = None,
+    overwrite: bool = False,
+):
+    """Standalone function for uploading a checkpoint file.
+
+    This function does not actually upload the checkpoint; it initiates the RemoteUploader's uploading of it
+    Args:
+        source_path (str): The path to the file to upload.
+        dest_dir (str): The directory/uri to upload the file to.
+        symlink_granularity (Optional[str]): The granularity to use for symlinking. One of 'file', 'dir', or None.
+            if None: no symlink uploaded
+            if 'file': command remoteuploader to wait until the file (specificied by source_path) is uploaded and then uploads a symlink pointing to the uploaded file
+            if 'dir': command remoteuploader  to wait until all files across all ranks are uploaded to dest_dir and then uploads a symlink
+                pointing to the remote directory (prefix in object_store terminology).
+        symlink_name (Optional[str]): The name to use for the symlink. Defaults to 'latest.symlink'.
+        async_upload (bool): If True, the uploads will be done asynchronously via the RemoteUploader and this function will return immediately.
+        state (Optional[State]): If async_upload is True, then state must be specified so that the remote_uploader can be
+            either extracted from state.callbacks or initialized and added to state.callbacks.
+        overwrite (bool): If allow overwrite existing remote checkpoint files
+    """
+    remote_uploader = RemoteUploader(remote_folder=dest_dir)
+    dest_path = ''
+    if source_path is not None:
+        _, _, dest_path = parse_uri(dest_dir)
+        remote_file_name = os.path.join(dest_path, os.path.basename(source_path))
+    else:
+        remote_file_name = None
+    all_remote_file_names = dist.all_gather_object([remote_file_name] if remote_file_name is not None else [])
+    remote_file_names = []
+    for filenames in all_remote_file_names:
+        remote_file_names += filenames
+
+    if source_path is not None:
+        assert remote_uploader is not None
+        assert remote_file_name is not None
+        remote_uploader.upload_file_async(
+            remote_file_name=remote_file_name,
+            file_path=pathlib.Path(source_path),
+            overwrite=overwrite,
+        )
+    symlink_remote_file_name = None
+    if dist.get_global_rank() == 0:
+        if symlink_name is not None:
+            symlink_remote_file_name = os.path.join(dest_path, symlink_name)
+            if symlink_granularity == 'file':
+                create_symlink_file(dest_path, symlink_name)
+            elif symlink_granularity == 'dir':
+                create_symlink_file(str(pathlib.Path(dest_path).parent), symlink_name)
+            elif symlink_granularity is not None:
+                raise ValueError(f'Unrecognized symlink granularity: {symlink_granularity}')
+    else:
+        symlink_remote_file_name = None
+        symlink_name = None
+
+    upload_callback = None
+    if state is not None:
+        for callback in state.callbacks:
+            if isinstance(callback, CheckpointUploadCallback):
+                upload_callback = callback
+                break
+    if upload_callback is None:
+        upload_callback = CheckpointUploadCallback(remote_uploader=remote_uploader)
+        if state is not None:
+            state.callbacks.append(upload_callback)
+    upload_callback.add_symlink_upload_task(
+        remote_file_names=remote_file_names,
+        symlink_file_name=symlink_name,
+        symlink_remote_file_name=symlink_remote_file_name,
+    )
+
+    if not async_upload:
+        upload_callback.wait()
diff --git a/composer/core/precision.py b/composer/core/precision.py
index bb91fc64d1..2723159261 100644
--- a/composer/core/precision.py
+++ b/composer/core/precision.py
@@ -9,6 +9,7 @@
 
 import torch
 
+from composer.devices import Device, DeviceCPU
 from composer.utils import StringEnum, is_xla_installed
 
 try:
@@ -17,7 +18,7 @@
 except ImportError:
     te_installed = False
 
-__all__ = ['Precision', 'get_precision_context']
+__all__ = ['Precision', 'get_precision_context', '_validate_precision']
 
 
 class Precision(StringEnum):
@@ -36,6 +37,12 @@ class Precision(StringEnum):
     AMP_FP8 = 'amp_fp8'
 
 
+def _validate_precision(precision: Precision, device: Device):
+    """Validate the precision and device combination."""
+    if isinstance(device, DeviceCPU) and precision != Precision.FP32:
+        raise ValueError(f'{precision} is not supported for CPU training.')
+
+
 @contextlib.contextmanager
 def get_precision_context(
     precision: Union[str, Precision],
diff --git a/composer/core/state.py b/composer/core/state.py
index a1bb14f0af..cbd7fc41db 100644
--- a/composer/core/state.py
+++ b/composer/core/state.py
@@ -528,6 +528,7 @@ def __init__(
         self.predict_timestamp = Timestamp()
         self._precision = Precision(precision)
         self._precision_config = precision_config
+        self.load_path: Optional[str] = None
 
         if optimizers is None:
             self._optimizers = []
@@ -546,6 +547,9 @@ def __init__(
         self.fsdp_config = parallelism_config.fsdp if parallelism_config is not None else None
         self.tp_config = parallelism_config.tp if parallelism_config is not None else None
 
+        self.automicrobatch_fsdp_hook_handles = []
+        self.fsdp_modules = {}
+
         self._validate_parallelism_configs()
 
         self.device_mesh: Optional[DeviceMesh] = _create_device_mesh(self.device, self.fsdp_config, self.tp_config)
@@ -636,9 +640,13 @@ def _validate_parallelism_configs(self):
             if error_message != '':
                 raise ValueError(error_message)
 
+        # Validate FSDP config parameters.
+        if self.fsdp_config is not None and self.fsdp_config.activation_cpu_offload and not self.fsdp_config.use_orig_params:
+            raise ValueError('activation_cpu_offload=True is not supported with use_orig_params=False.')
+
         # Validate FSDP state dict type
-        if self.fsdp_state_dict_type not in [None, 'full', 'sharded']:
-            if self.fsdp_state_dict_type == 'local':
+        if self.fsdp_config is not None and self.fsdp_config.state_dict_type not in [None, 'full', 'sharded']:
+            if self.fsdp_config.state_dict_type == 'local':
                 raise ValueError(
                     'Composer and PyTorch no longer support saving or loading local state dicts. '
                     'To upgrade an older checkpoint, use Composer version 0.18.1 and export as '
@@ -646,7 +654,7 @@ def _validate_parallelism_configs(self):
                 )
             raise ValueError(
                 f'fsdp_state_dict_type must be one of [None, "full", "sharded"], but got '
-                f'{self.fsdp_state_dict_type}',
+                f'{self.fsdp_config.state_dict_type}',
             )
         if self.fsdp_sharded_state_dict_enabled and self.save_metrics:
             # Sharded state dict breaks in many different ways with torchmetrics, due to both sharding
@@ -1386,7 +1394,7 @@ def load_model_state(
             with reproducibility.seed_context(self.rank_zero_seed):
                 from composer.distributed import prepare_fsdp_module
 
-                prepare_fsdp_module(
+                self.automicrobatch_fsdp_hook_handles, self.fsdp_modules = prepare_fsdp_module(
                     self.model,
                     self.optimizers,
                     self.fsdp_config,
diff --git a/composer/distributed/dist_strategy.py b/composer/distributed/dist_strategy.py
index be81652881..f7adc79428 100644
--- a/composer/distributed/dist_strategy.py
+++ b/composer/distributed/dist_strategy.py
@@ -7,7 +7,7 @@
 import logging
 import warnings
 from contextlib import contextmanager, nullcontext
-from typing import Any, Callable, ContextManager, Iterator, Optional, Sequence, Union, cast
+from typing import Any, Callable, ContextManager, Iterator, Optional, Sequence, Tuple, Union, cast
 
 import torch
 from packaging import version
@@ -22,7 +22,8 @@
 from torchmetrics import Metric, MetricCollection
 
 from composer.core import Precision, State
-from composer.devices import Device
+from composer.core.precision import _validate_precision
+from composer.devices import Device, DeviceGPU
 from composer.distributed.meta_safe_apply import meta_safe_apply
 from composer.distributed.mosaic_parallelism import (
     BACKWARD_PREFETCH_MAP,
@@ -31,7 +32,7 @@
     get_mixed_precision,
     set_custom_fsdp_module_kwargs,
 )
-from composer.utils import FSDPConfig, StringEnum, TPConfig, dist, ensure_tuple
+from composer.utils import FSDPConfig, StringEnum, TPConfig, dist, ensure_tuple, get_device
 
 __all__ = ['DDPSyncStrategy', 'ddp_sync_context', 'prepare_ddp_module', 'prepare_fsdp_module', 'prepare_tp_module']
 
@@ -171,9 +172,14 @@ def _recreate_fsdp_param_groups_from_unwrapped_opt_info(
     for fsdp_name, param in fsdp_wrapped_named_params:
 
         unwrapped_name = clean_tensor_name(fsdp_name)
-        # need to have a 1:1 mapping between a fsdp param name and the non-wrapped vanilla param name
-        retrieved_group_num = non_wrapped_param_names_to_group_num[unwrapped_name]
-        group_num_to_optimizer_info[retrieved_group_num]['params'].append(param)
+
+        # Since we are iterating over all model.named_parameters() after fsdp wrapping, we need to check
+        # if the parameter was included in the optimizer param_group pre fsdp wrapping, in order to support
+        # passing a subset of model params in the optimizer
+        if unwrapped_name in non_wrapped_param_names_to_group_num:
+            # Need to have a 1:1 mapping between a fsdp param name and the non-wrapped vanilla param name
+            retrieved_group_num = non_wrapped_param_names_to_group_num[unwrapped_name]
+            group_num_to_optimizer_info[retrieved_group_num]['params'].append(param)
 
     # return sorted optimizer info groups
     return [group_num_to_optimizer_info[num] for num in sorted(group_num_to_optimizer_info.keys())]
@@ -181,9 +187,23 @@ def _recreate_fsdp_param_groups_from_unwrapped_opt_info(
 
 def prepare_tp_module(
     model: torch.nn.Module,
+    optimizers: Optional[Union[torch.optim.Optimizer, Sequence[torch.optim.Optimizer]]],
     tp_config: TPConfig,
 ) -> None:
     """Prepare a module (assumed ComposerModel) for use with tensor parallel."""
+    optimizers_tuple = ensure_tuple(optimizers)
+    if len(optimizers_tuple) != 1:
+        raise NotImplementedError(f'Only one optimizer is supported; found {len(optimizers_tuple)} optimizers')
+
+    optim = optimizers_tuple[0]
+    if len(optim.param_groups) > 1:
+        raise RuntimeError('Multiple optimizer groups are not supported with tensor parallelism.',)
+
+    if len(optim.param_groups[0]['params']) != len(list(model.parameters())):
+        raise ValueError(
+            'Passing in a subset of model parameters to the optimizer is not supported with tensor parallelism.',
+        )
+
     from torch.distributed.tensor.parallel import parallelize_module
 
     device_mesh = tp_config.device_mesh
@@ -199,11 +219,11 @@ def prepare_fsdp_module(
     model: torch.nn.Module,
     optimizers: Optional[Union[torch.optim.Optimizer, Sequence[torch.optim.Optimizer]]],
     fsdp_config: FSDPConfig,
-    precision: Precision,
-    device: Device,
-    auto_microbatching: bool,
+    precision: Optional[Union[str, Precision]] = None,
+    device: Optional[Union[str, Device]] = None,
+    auto_microbatching: bool = False,
     te_rng_seed: int = 1234,
-) -> None:
+) -> Tuple[list, dict]:
     """Prepare a module (assumed ComposerModel) and optimizer for use with :class:`torch.distributed.fsdp.FullyShardedDataParallel`.
 
     Args:
@@ -211,10 +231,18 @@ def prepare_fsdp_module(
         optimizers (torch.optim.Optimizer | Sequence[torch.optim.Optimizer], optional): The optimizer for `model`, assumed to have a single param group := model.parameters().
         fsdp_config (FSDPConfig): The FSDP config.
         precision: (Precision): The precision being used by the Trainer, used to fill in defaults for FSDP `mixed_precision` settings.
-        device (Device): The device being used by the Trainer.
+        device:  The device being used by the Trainer.
         auto_microbatching (bool, optional): Whether or not auto microbatching is enabled.
         te_rng_seed(int): The seed to use for the Transformer Engine activation checkpointing RNG. Defaults to 1234.
     """
+    device = get_device(device)
+
+    if precision is None:
+        precision = Precision.AMP_FP16 if isinstance(device, DeviceGPU) else Precision.FP32
+    elif isinstance(precision, str):
+        precision = Precision(precision)
+    _validate_precision(precision, device)
+
     # Check sync_module_states is True for mixed initialization or HSDP
     if fsdp_config.sync_module_states == False:
         rank_on_meta = 1 if next(model.parameters()).device.type == 'meta' else 0
@@ -230,6 +258,9 @@ def prepare_fsdp_module(
                 'some weights may be randomly initialized when loading a checkpoint.',
             )
 
+    # Handles of FSDP sync hooks if automicrobatching is on
+    hook_handles = []
+
     # Check if other ranks OOMed after forward/backward pass when using auto microbatching. This
     # may happen when close to memory limit or with uneven memory usage across ranks. Since we
     # need to do this before the model weights are gathered for the next FSDP block, we wrap every
@@ -247,11 +278,10 @@ def sync_hook(*args):
             raise RuntimeError('CUDA out of memory encountered on a different rank')
 
     # Necessary variables for optimizers with multiple param groups in FSDP
-    num_param_groups = None
     param_name_to_group_num = None
-    group_num_to_param_group_info = None
+    group_num_to_opt_group_info = None
+    single_param_group_opt_info = None
 
-    optimizer_specific_info = None
     if optimizers:
         optimizers_tuple = ensure_tuple(optimizers)
         if len(optimizers_tuple) != 1:
@@ -261,38 +291,41 @@ def sync_hook(*args):
         # that will be recreated at the end of prepare_fsdp_module
         optim = optimizers_tuple[0]
 
-        num_param_groups = len(optim.param_groups)
-        if num_param_groups > 1:
-            if not fsdp_config.use_orig_params:
-                raise RuntimeError(
-                    'Multiple optimizer groups with FSDP are only supported with '
-                    'use_orig_params=True.',
-                )
-            # optimizer.param_groups do not contain parameter names which are needed
-            # to keep track of the different parameters in each group
-            # so we use the pointers between model.parameters() and model.named_parameters()
-            # to get the names of the parameters within optimizer.param_groups
-            param_pointer_to_param_name = {id(p): n for n, p in model.named_parameters()}
+        # Simplest case - single param group & all model params stored in optimizer
+        if len(optim.param_groups) == 1 and len(optim.param_groups[0]['params']) == len(list(model.parameters())):
+            single_param_group_opt_info = {k: v for k, v in optim.param_groups[0].items() if k != 'params'}
+        elif fsdp_config.use_orig_params:
+            # this code block stores information about param groups pre-fsdp wrapping in order to recreate them post-wrapping
+            # to do so, it relies on the ptrs of the model.parameters() in a model and the names of the params
+            # for this to work, use_orig_params=True, as we need the names of the params post-wrapping
+            # TP is not supported, as the underlying parameters in the model differ from the params in the param groups after being dtensorified
+
+            ptr_to_param_name = {id(p): n for n, p in model.named_parameters()}
             param_name_to_group_num = {}
-            group_num_to_param_group_info = {}
+            group_num_to_opt_group_info = {}
             for group_num in range(len(optim.param_groups)):
                 # Need to in-line to avoid a reference which causes FSDP to allocate extra GPU memory
                 # group = optim.param_groups[group_num]
                 for param_num in range(len(optim.param_groups[group_num]['params'])):
-                    # Need to in-line to avoid a reference which causes FSDP to allocate extra GPU memory
-                    # param = optim.param_groups[group_num]['params'][param_num]
-                    param_name_to_group_num[param_pointer_to_param_name[id(
-                        optim.param_groups[group_num]['params'][param_num],
-                    )]] = group_num
+                    param_ptr = id(optim.param_groups[group_num]['params'][param_num])
+                    if param_ptr not in ptr_to_param_name:
+                        raise ValueError('The same model must be passed to the optimizer and trainer.')
+                    param_name_to_group_num[ptr_to_param_name[param_ptr]] = group_num
 
                 # this includes optimizer-specific values like lr, eps
                 # this will be used as the kwargs for the optim param groups later
                 optimizer_specific_group_info = {
                     k: v for k, v in optim.param_groups[group_num].items() if k != 'params'
                 }
-                group_num_to_param_group_info[group_num] = optimizer_specific_group_info
+                group_num_to_opt_group_info[group_num] = optimizer_specific_group_info
         else:
-            optimizer_specific_info = {k: v for k, v in optim.param_groups[0].items() if k != 'params'}
+            if len(optim.param_groups) > 1:
+                raise RuntimeError('Multiple optimizer groups with FSDP are not supported with use_orig_params=False.',)
+
+            if len(optim.param_groups[0]['params']) != len(list(model.parameters())):
+                raise ValueError(
+                    'Passing in a subset of model parameters to the optimizer is not supported with use_orig_params=False.',
+                )
 
         optim.param_groups.clear()
         optim.state.clear()
@@ -352,6 +385,7 @@ def sync_hook(*args):
     sharded_ckpt_prefix_dir = fsdp_config.sharded_ckpt_prefix_dir
     use_orig_params = fsdp_config.use_orig_params
 
+    fsdp_obj_named_modules = {}
     # We choose to not wrap the ComposerModel directly, but instead wrap any submodules like `ComposerModel.model`
     # This makes it safer to call ComposerModel-specific functions like 'eval_forward' that
     # may make calls to sharded submodules. If we only wrap the submodules, then any call that ComposerModel makes
@@ -512,9 +546,6 @@ def lambda_fn(module: torch.nn.Module) -> Union[bool, dict]:
                         ret = obj.fsdp_wrap_fn(module)
                         if isinstance(ret, dict):
                             ret = set_custom_fsdp_module_kwargs(ret, process_group_cache)
-                    if ret and auto_microbatching:
-                        module.register_forward_hook(sync_hook)
-                        module.register_full_backward_hook(sync_hook)
                     return ret
 
                 _auto_wrap_policy = CustomPolicy(lambda_fn)
@@ -531,9 +562,6 @@ def __auto_wrap_policy(module: torch.nn.Module, recurse: bool, nonwrapped_numel:
                     elif hasattr(obj, 'fsdp_wrap_fn') and isinstance(obj.fsdp_wrap_fn, Callable):
                         should_be_wrapped = obj.fsdp_wrap_fn(module)
 
-                    if should_be_wrapped and auto_microbatching:
-                        module.register_forward_hook(sync_hook)
-                        module.register_full_backward_hook(sync_hook)
                     return should_be_wrapped
 
                 def _auto_wrap_policy_new(module: torch.nn.Module, recurse: bool, nonwrapped_numel: int) -> bool:
@@ -567,6 +595,22 @@ def _auto_wrap_policy_new(module: torch.nn.Module, recurse: bool, nonwrapped_num
                 log.info(f'Calling prepare_te_modules_for_fsdp to enable TE weights sharding')
                 prepare_te_modules_for_fsdp(fsdp_obj)
 
+            # The following sync hooks are added to prevent FSDP deadlocks that are caused when some ranks OOM
+            # and other ranks do not OOM, leading to OOMing ranks calling all_reduce to wait on the non-OOMing
+            # ranks and the non-OOMing ranks calling all_gatherbase to continue with FSDP training:
+            #
+            #   forward_pre_hook: before forwards of FSDP modules
+            #   full_backward_pre_hook: before backwards of FSDP modules
+            #   full_backward_hook: before a prefetched unshard called by FSDP's `post_backward_reshard`
+            if auto_microbatching:
+                for _, module in fsdp_obj.named_modules():
+                    if isinstance(module, FullyShardedDataParallel):
+                        hook_handles.append(module.register_forward_pre_hook(sync_hook, prepend=True))
+                        hook_handles.append(module.register_full_backward_pre_hook(sync_hook, prepend=True))
+                    else:
+                        hook_handles.append(module.register_full_backward_hook(sync_hook))
+                fsdp_obj_named_modules.update(dict(fsdp_obj.named_modules()))
+
             if hasattr(fsdp_obj, '_exec_order_data'):
                 if hasattr(fsdp_obj._exec_order_data, '_forward_prefetch_limit'):
                     fsdp_obj._exec_order_data._forward_prefetch_limit = fsdp_config.forward_prefetch_limit
@@ -711,19 +755,19 @@ def _check_fn(module: torch.nn.Module) -> bool:
         optim = ensure_tuple(optimizers)[0]
         optim.param_groups.clear()
 
-        assert num_param_groups is not None
-        if num_param_groups > 1:
+        if single_param_group_opt_info is not None:
+            single_param_group_opt_info.update({'params': list(model.parameters())})
+            optim.add_param_group(single_param_group_opt_info)
+        elif fsdp_config.use_orig_params:
             assert param_name_to_group_num is not None
-            assert group_num_to_param_group_info is not None
+            assert group_num_to_opt_group_info is not None
 
             param_groups = _recreate_fsdp_param_groups_from_unwrapped_opt_info(
                 model.named_parameters(),
                 param_name_to_group_num,
-                group_num_to_param_group_info,
+                group_num_to_opt_group_info,
             )
             for param_group in param_groups:
                 optim.add_param_group(param_group)
-        else:
-            assert optimizer_specific_info is not None
-            optimizer_specific_info.update({'params': list(model.parameters())})
-            optim.add_param_group(optimizer_specific_info)
+
+    return hook_handles, fsdp_obj_named_modules
diff --git a/composer/loggers/mlflow_logger.py b/composer/loggers/mlflow_logger.py
index c90f167b82..7a8acdd511 100644
--- a/composer/loggers/mlflow_logger.py
+++ b/composer/loggers/mlflow_logger.py
@@ -7,9 +7,12 @@
 
 import fnmatch
 import logging
+import multiprocessing
 import os
 import pathlib
 import posixpath
+import signal
+import sys
 import textwrap
 import time
 import warnings
@@ -33,6 +36,52 @@
 DEFAULT_MLFLOW_EXPERIMENT_NAME = 'my-mlflow-experiment'
 
 
+class MlflowMonitorProcess(multiprocessing.Process):
+
+    def __init__(self, main_pid, mlflow_run_id, mlflow_tracking_uri):
+        super().__init__()
+        self.main_pid = main_pid
+        self.mlflow_run_id = mlflow_run_id
+        self.mlflow_tracking_uri = mlflow_tracking_uri
+        self.exit_event = multiprocessing.Event()
+        self.crash_event = multiprocessing.Event()
+
+    def handle_sigterm(self, signum, frame):
+        from mlflow import MlflowClient
+        client = MlflowClient(self.mlflow_tracking_uri)
+        if client.get_run(self.mlflow_run_id).info.status == 'RUNNING':
+            # Set the run status as KILLED if SIGTERM is received while the MLflow run is still
+            # in status RUNNING.
+            client.set_terminated(self.mlflow_run_id, status='KILLED')
+
+    def run(self):
+        from mlflow import MlflowClient
+
+        os.setsid()
+        # Register the signal handler in the child process
+        signal.signal(signal.SIGTERM, self.handle_sigterm)
+
+        while not self.exit_event.wait(10):
+            try:
+                # Signal 0 does not kill the process but performs error checking
+                os.kill(self.main_pid, 0)
+            except OSError:
+                client = MlflowClient(self.mlflow_tracking_uri)
+                client.set_terminated(self.mlflow_run_id, status='FAILED')
+                break
+
+        if self.crash_event.is_set():
+            client = MlflowClient(self.mlflow_tracking_uri)
+            client.set_terminated(self.mlflow_run_id, status='FAILED')
+
+    def stop(self):
+        self.exit_event.set()
+
+    def crash(self):
+        self.crash_event.set()
+        self.exit_event.set()
+
+
 class MLFlowLogger(LoggerDestination):
     """Log to `MLflow <https://www.mlflow.org/docs/latest/index.html>`_.
 
@@ -121,7 +170,7 @@ def __init__(
         self.resume = resume
 
         if logging_buffer_seconds:
-            os.environ['MLFLOW_ASYNC_LOGGING_BUFFERING_SECONDS'] = str(logging_buffer_seconds)
+            os.environ['MLFLOW_ASYNC_LOGGING_BUFFERING_SECONDS'] = str(logging_buffer_seconds,)
 
         if log_system_metrics:
             # Set system metrics sampling interval and samples before logging so that system metrics
@@ -172,7 +221,11 @@ def __init__(
                         conda_channel='conda-forge',
                     ) from e
                 databricks_username = WorkspaceClient().current_user.me().user_name or ''
-                self.experiment_name = os.path.join('/Users', databricks_username, self.experiment_name.strip('/'))
+                self.experiment_name = os.path.join(
+                    '/Users',
+                    databricks_username,
+                    self.experiment_name.strip('/'),
+                )
 
             self._mlflow_client = MlflowClient(self.tracking_uri)
             # Set experiment
@@ -187,7 +240,7 @@ def __init__(
                 if exp_from_name is not None:
                     self._experiment_id = exp_from_name.experiment_id
                 else:
-                    self._experiment_id = (self._mlflow_client.create_experiment(name=self.experiment_name))
+                    self._experiment_id = self._mlflow_client.create_experiment(name=self.experiment_name)
 
     def _start_mlflow_run(self, state):
         import mlflow
@@ -239,6 +292,19 @@ def _start_mlflow_run(self, state):
             tags=self.tags,
             log_system_metrics=self.log_system_metrics,
         )
+        if self.tracking_uri == 'databricks':
+            # Start a background process to monitor the job to report the job status to MLflow.
+            self.monitor_process = MlflowMonitorProcess(
+                os.getpid(),
+                self._run_id,
+                self.tracking_uri,
+            )
+            self.monitor_process.start()
+
+    def _global_exception_handler(self, exc_type, exc_value, exc_traceback):
+        """Catch global exception."""
+        self._global_exception_occurred += 1
+        sys.__excepthook__(exc_type, exc_value, exc_traceback)
 
     def init(self, state: State, logger: Logger) -> None:
         del logger  # unused
@@ -246,6 +312,11 @@ def init(self, state: State, logger: Logger) -> None:
         if self.run_name is None:
             self.run_name = state.run_name
 
+        if hasattr(state, 'device'):
+            self._global_exception_occurred = state.device.tensor_to_device(torch.tensor([0], dtype=torch.uint8),)
+        else:
+            self._global_exception_occurred = 0
+
         # Store the Composer run name in the MLFlow run tags so it can be retrieved for autoresume
         self.tags['run_name'] = os.environ.get('RUN_NAME', state.run_name)
 
@@ -253,6 +324,8 @@ def init(self, state: State, logger: Logger) -> None:
         if not self._rank_zero_only:
             self.run_name += f'-rank{dist.get_global_rank()}'
 
+        # Register the global exception handler so that uncaught exception is tracked.
+        sys.excepthook = self._global_exception_handler
         # Start run
         if self._enabled:
             self._start_mlflow_run(state)
@@ -265,7 +338,10 @@ def init(self, state: State, logger: Logger) -> None:
             self._experiment_id, self._run_id = mlflow_ids_list
 
     def after_load(self, state: State, logger: Logger) -> None:
-        logger.log_hyperparameters({'mlflow_experiment_id': self._experiment_id, 'mlflow_run_id': self._run_id})
+        logger.log_hyperparameters({
+            'mlflow_experiment_id': self._experiment_id,
+            'mlflow_run_id': self._run_id,
+        })
         self.run_url = posixpath.join(
             os.environ.get('DATABRICKS_HOST', ''),
             'ml',
@@ -305,6 +381,7 @@ def rename(self, key: str):
 
     def log_metrics(self, metrics: dict[str, Any], step: Optional[int] = None) -> None:
         from mlflow import log_metrics
+
         if self._enabled:
             # Convert all metrics to floats to placate mlflow.
             metrics = {
@@ -357,6 +434,7 @@ def register_model(
             full_name = f'{self.model_registry_prefix}.{name}' if len(self.model_registry_prefix) > 0 else name
 
             import mlflow
+
             return mlflow.register_model(
                 model_uri=model_uri,
                 name=full_name,
@@ -378,6 +456,7 @@ def save_model(self, flavor: Literal['transformers', 'peft'], **kwargs):
         """
         if self._enabled:
             import mlflow
+
             if flavor == 'transformers':
                 mlflow.transformers.save_model(**kwargs)
             elif flavor == 'peft':
@@ -443,6 +522,7 @@ def log_model(self, flavor: Literal['transformers'], **kwargs):
         """
         if self._enabled:
             import mlflow
+
             if flavor == 'transformers':
                 mlflow.transformers.log_model(**kwargs)
             else:
@@ -528,21 +608,49 @@ def log_images(
 
     def post_close(self):
         if self._enabled:
+            if hasattr(self, 'monitor_process'):
+                # Check if there is an uncaught exception, which means `post_close()` is triggered
+                # due to program crash.
+                if isinstance(self._global_exception_occurred, torch.Tensor):
+                    finish_with_exception = (self._global_exception_occurred == 1).item()
+                else:
+                    finish_with_exception = (self._global_exception_occurred == 1)
+                if finish_with_exception:
+                    self.monitor_process.crash()
+                    return
+
+                # Stop the monitor process since it's entering the cleanup phase.
+                self.monitor_process.stop()
+
             import mlflow
 
             assert isinstance(self._run_id, str)
+
             mlflow.flush_async_logging()
-            self._mlflow_client.set_terminated(self._run_id)
+            exc_tpe, exc_info, tb = sys.exc_info()
+            if (exc_tpe, exc_info, tb) == (None, None, None):
+                current_status = self._mlflow_client.get_run(self._run_id).info.status
+                if current_status == 'RUNNING':
+                    self._mlflow_client.set_terminated(self._run_id, status='FINISHED')
+            else:
+                # Record there was an error
+                self._mlflow_client.set_terminated(self._run_id, status='FAILED')
+
             mlflow.end_run()
+            if hasattr(self, 'monitor_process'):
+                self.monitor_process.join()
 
 
-def _convert_to_mlflow_image(image: Union[np.ndarray, torch.Tensor], channels_last: bool) -> np.ndarray:
+def _convert_to_mlflow_image(
+    image: Union[np.ndarray, torch.Tensor],
+    channels_last: bool,
+) -> np.ndarray:
     if isinstance(image, torch.Tensor):
         image = image.data.cpu().numpy()
 
     # Error out for empty arrays or weird arrays of dimension 0.
     if np.any(np.equal(image.shape, 0)):
-        raise ValueError(f'Got an image (shape {image.shape}) with at least one dimension being 0! ')
+        raise ValueError(f'Got an image (shape {image.shape}) with at least one dimension being 0!')
 
     # Squeeze any singleton dimensions and then add them back in if image dimension
     # less than 3.
diff --git a/composer/loggers/mosaicml_logger.py b/composer/loggers/mosaicml_logger.py
index 2de7243d31..1fcb612ed5 100644
--- a/composer/loggers/mosaicml_logger.py
+++ b/composer/loggers/mosaicml_logger.py
@@ -98,6 +98,11 @@ def log_metrics(self, metrics: dict[str, Any], step: Optional[int] = None) -> No
         self.log_metadata(metrics)
 
     def after_load(self, state: State, logger: Logger) -> None:
+        # If model was resumed from checkpoint, log the checkpoint path
+        if int(state.timestamp.batch) > 0 and state.load_path is not None:
+            log.debug(f'Logging checkpoint path to metadata: {state.load_path}')
+            self.log_metadata({'checkpoint_resumed_path': state.load_path})
+
         # Log model data downloaded and initialized for run events
         log.debug(f'Logging model initialized time to metadata')
         self.log_metadata({'model_initialized_time': time.time()})
diff --git a/composer/loggers/wandb_logger.py b/composer/loggers/wandb_logger.py
index 4a2afe5c84..d76ee1fbac 100644
--- a/composer/loggers/wandb_logger.py
+++ b/composer/loggers/wandb_logger.py
@@ -200,7 +200,11 @@ def init(self, state: State, logger: Logger) -> None:
         if self._enabled:
             wandb.init(**self._init_kwargs)
             assert wandb.run is not None, 'The wandb run is set after init'
-            entity_and_project = [str(wandb.run.entity), str(wandb.run.project)]
+            if hasattr(wandb.run, 'entity') and hasattr(wandb.run, 'project'):
+                entity_and_project = [str(wandb.run.entity), str(wandb.run.project)]
+            else:
+                # Run does not have attribtues if wandb is in disabled mode, so we must mock it
+                entity_and_project = ['disabled', 'disabled']
             self.run_dir = wandb.run.dir
             self.run_url = wandb.run.get_url()
             atexit.register(self._set_is_in_atexit)
diff --git a/composer/trainer/_patch_pytorch.py b/composer/trainer/_patch_pytorch.py
index 3f19df7d2a..2c27118090 100644
--- a/composer/trainer/_patch_pytorch.py
+++ b/composer/trainer/_patch_pytorch.py
@@ -31,8 +31,19 @@
 from torch.distributed.fsdp._fsdp_extensions import _ext_pre_load_state_dict_transform
 from torch.distributed.utils import _replace_by_prefix
 
+from composer.utils import dist
+
 log = logging.getLogger(__name__)
 
+def patch_unshard_for_automicrobatching(auto_microbatch_size_found=False):
+    """Monkey patches sync hook into unshard when searching during automicrobatching."""
+    if version.parse(torch.__version__) >= version.parse('2.3.1'):
+        from torch.distributed.fsdp._flat_param import FlatParamHandle
+        if auto_microbatch_size_found:
+            global original_unshard
+            FlatParamHandle.unshard = (original_unshard)
+        else:
+            FlatParamHandle.unshard = (unshard_with_sync)
 
 def patch_pytorch():
     """Monkey patches pytorch functions based on pytorch version."""
@@ -1043,3 +1054,49 @@ def device_mesh__getitem__(
 
             submesh = _mesh_resources.create_child_mesh(self, mesh_dim_names)
             return submesh
+
+
+    # Save original FlatParamHandle.unshard to revert back to when dropping automicrobatching hooks
+    from torch.distributed.fsdp._flat_param import FlatParamHandle
+    original_unshard = FlatParamHandle.unshard
+
+    @no_type_check
+    def unshard_with_sync(self):
+        """Run the unshard logic, but with a sync after a :meth:`_alloc_padded_unsharded_flat_param`.
+
+        This prevents deadlocks when some ranks OOM after the alloc call and others do not.
+        This is a patched method from pytorch, meant to be called when automicrobatching
+        turns on hooks in its search process for the optimal non-OOMing microbatch size.
+        This includes all-gathering the flat parameter
+        and switching to using the unsharded flat parameter. If the handle does
+        not need unsharding, then this only switches to using the unsharded
+        flat parameter. For ``NO_SHARD``, this is a no-op.
+        If FSDP is in :meth:`summon_full_params` and the handle uses parameter
+        mixed precision, then the parameter is forced to full precision.
+        """
+        if not self.needs_unshard():
+            # Even when not needing an unshard, we should switch to using
+            # the unsharded flat parameter
+            unsharded_flat_param = (
+                self._get_padded_unsharded_flat_param()
+                if self.uses_sharded_strategy
+                else self.flat_param
+            )
+            self._use_unsharded_flat_param(unsharded_flat_param)
+            return
+        unsharded_flat_param = self._alloc_padded_unsharded_flat_param()
+
+        # Check if any other rank hit an OOM
+        found_cuda_oom_tensor = torch.tensor([0], dtype=torch.uint8).to(self.device, non_blocking=True)
+
+        dist.all_reduce(found_cuda_oom_tensor, reduce_operation='MAX')
+        found_cuda_oom = found_cuda_oom_tensor.item()
+        # Signal current rank is still in batch
+        all_ranks_finished_tensor = torch.tensor([0], dtype=torch.uint8).to(self.device, non_blocking=True)
+
+        dist.all_reduce(all_ranks_finished_tensor, reduce_operation='MIN')
+
+        if found_cuda_oom == 1:
+            raise RuntimeError('CUDA out of memory encountered on a different rank')
+        padded_unsharded_flat_param = self._all_gather_flat_param(unsharded_flat_param)
+        self._use_unsharded_flat_param(padded_unsharded_flat_param)
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
index b62b3d3e58..b2f829ca10 100644
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -24,6 +24,7 @@
     Any,
     Callable,
     ContextManager,
+    Dict,
     Iterable,
     Mapping,
     Optional,
@@ -75,6 +76,7 @@
     ensure_time,
     get_precision_context,
 )
+from composer.core.precision import _validate_precision
 from composer.devices import Device, DeviceCPU, DeviceGPU, DeviceMPS, DeviceTPU
 from composer.distributed import (
     DDPSyncStrategy,
@@ -99,7 +101,7 @@
 from composer.models import ComposerModel
 from composer.optim import ComposerScheduler, DecoupledSGDW, compile_composer_scheduler
 from composer.profiler import Profiler
-from composer.trainer._patch_pytorch import patch_pytorch
+from composer.trainer._patch_pytorch import patch_pytorch, patch_unshard_for_automicrobatching
 from composer.trainer._scale_schedule import scale_pytorch_scheduler
 from composer.trainer._scaler import ClosureGradScaler
 from composer.utils import (
@@ -145,6 +147,8 @@
 # syntax to shorten the Scheduler type annotations
 Scheduler = Union[ComposerScheduler, LRScheduler]
 
+OOM_FOUND_ON_OTHER_RANK = 'CUDA out of memory encountered on a different rank'
+
 
 def _raise_missing_argument_exception(arg_name: str):
     raise ValueError((
@@ -198,11 +202,6 @@ def _filter_metrics(metrics: dict[str, Metric], metric_names: Optional[list[str]
     return filtered_metrics
 
 
-def _validate_precision(precision: Precision, device: Device):
-    if isinstance(device, DeviceCPU) and precision != Precision.FP32:
-        raise ValueError(f'{precision} is not supported for CPU training.')
-
-
 def _compile_schedulers(
     schedulers: Optional[Union[Scheduler, Sequence[Scheduler]]],
     state: State,
@@ -338,6 +337,25 @@ def _fsdp_reshard_and_cleanup(model: torch.nn.Module):
                 _post_backward_final_callback(module, module)
 
 
+def _clear_incomplete_train_states(state: State):
+    """Manually clear gradients when automicrobatching reruns a batch.
+
+    Before automicrobatching tries a lower microbatch size, clear the
+    training states and memory of the previous run of the batch to reset the memory to
+    before the batch was run.
+    """
+    if hasattr(state, 'outputs'):
+        del state.outputs
+    if hasattr(state, 'loss'):
+        del state.loss
+    for optimizer in state.optimizers:
+        optimizer.zero_grad(set_to_none=True)
+    if state.scaler is not None:
+        state.scaler._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
+    _fsdp_reshard_and_cleanup(state.model)
+    torch.cuda.empty_cache()
+
+
 def _adjust_device_train_microbatch_size(state: State):
     """Adjust device_train_microbatch_size if we encounter OOM.
 
@@ -349,7 +367,7 @@ def _adjust_device_train_microbatch_size(state: State):
     assert state.device_train_microbatch_size is not None
     if state.device_train_microbatch_size == 1:
         raise RuntimeError((
-            'CUDA out of memory. The train loop failed with an internal microbatch of size 1.'
+            'CUDA out of memory or excessive memory allocation retries detected. The train loop failed with an internal microbatch of size 1.'
             'The GPU does not have enough memory to process even 1 sample during train.'
         ))
     else:
@@ -357,21 +375,12 @@ def _adjust_device_train_microbatch_size(state: State):
         state.device_train_microbatch_size = max(int(original_microbatch_size / 2), 1)
         warnings.warn(
             RuntimeWarning(
-                'CUDA out of memory detected. Train microbatch size will be decreased from '
+                'CUDA out of memory or excessive memory allocation retries detected. Train microbatch size will be decreased from '
                 f'{original_microbatch_size} -> {state.device_train_microbatch_size}.',
             ),
         )
     # Clear gradients in case failure happened during backwards pass
-    if hasattr(state, 'outputs'):
-        del state.outputs
-    if hasattr(state, 'loss'):
-        del state.loss
-    for optimizer in state.optimizers:
-        optimizer.zero_grad(set_to_none=True)
-    if state.scaler is not None:
-        state.scaler._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)
-    _fsdp_reshard_and_cleanup(state.model)
-    torch.cuda.empty_cache()
+    _clear_incomplete_train_states(state)
 
 
 def _adjust_device_eval_microbatch_size(evaluator: Evaluator):
@@ -400,6 +409,75 @@ def _adjust_device_eval_microbatch_size(evaluator: Evaluator):
     torch.cuda.empty_cache()
 
 
+def _update_num_consecutive_thrashes(state: State, num_consecutive_thrashes: int, num_alloc_retries: int):
+    """Update the number of consecutive batches where we experienced alloc retries.
+
+    Consecutive alloc retries in GPU memory usually indicate thrashing, where GPU memory usage is so close
+    to the memory limit that it hinders throughput.
+    """
+    # Check for alloc retries between batches
+    stats = torch.cuda.memory_stats()
+    cur_num_alloc_retries = stats['num_alloc_retries']
+
+    if cur_num_alloc_retries - num_alloc_retries > 0:
+        alloc_retry_this_batch = 1
+        log.info('Found new alloc retries this batch: ' + str(num_alloc_retries) + ' to ' + str(cur_num_alloc_retries))
+    else:
+        alloc_retry_this_batch = 0
+
+    # Propagate across all ranks if any rank had alloc retries this batch
+    alloc_retry_tensor = state.device.tensor_to_device(torch.tensor([alloc_retry_this_batch], dtype=torch.uint8),)
+    dist.all_reduce(alloc_retry_tensor, reduce_operation='MAX')
+    alloc_retry_this_batch = alloc_retry_tensor.item() == 1
+    if alloc_retry_this_batch:
+        num_consecutive_thrashes += 1
+    else:
+        num_consecutive_thrashes = 0
+    return num_consecutive_thrashes
+
+
+def _create_sync_hook(state: State):
+    """Check if other ranks OOMed after forward/backward pass when using auto microbatching.
+
+    This may happen when close to memory limit or with uneven memory usage across ranks. Since we
+    need to do this before the model weights are gathered for the next FSDP block, we wrap every
+    FSPD block with a hook that checks if any other rank OOMed.
+
+    This wrapper method is needed because PyTorch FSDP doesn't take `state` as an argument in hooks
+    that are registered using methods such as `register_forward_pre_hook`.
+    """
+
+    def sync_hook(*args):
+        # Check if any other rank hit an OOM
+        found_cuda_oom_tensor = state.device.tensor_to_device(torch.tensor([0], dtype=torch.uint8))
+        dist.all_reduce(found_cuda_oom_tensor, reduce_operation='MAX')
+        found_cuda_oom = found_cuda_oom_tensor.item()
+        # Signal current rank is still in batch
+        all_ranks_finished_tensor = state.device.tensor_to_device(torch.tensor([0], dtype=torch.uint8))
+        dist.all_reduce(all_ranks_finished_tensor, reduce_operation='MIN')
+
+        if found_cuda_oom == 1:
+            raise RuntimeError()
+
+    return sync_hook
+
+
+def _readd_fsdp_sync_hooks(fsdp_modules: Dict[str, torch.nn.Module], sync_hook):
+    """Readds previously removed sync hooks back to FSDP modules.
+
+    Called when preparing to search for or searching for new microbatch size during automicrobatching.
+    """
+    automicrobatch_fsdp_hook_handles = []
+    patch_unshard_for_automicrobatching(auto_microbatch_size_found=False)
+    for module in fsdp_modules.values():
+        if isinstance(module, FullyShardedDataParallel):
+            automicrobatch_fsdp_hook_handles.append(module.register_forward_pre_hook(sync_hook, prepend=True))
+            automicrobatch_fsdp_hook_handles.append(module.register_full_backward_pre_hook(sync_hook, prepend=True))
+        else:
+            automicrobatch_fsdp_hook_handles.append(module.register_full_backward_hook(sync_hook))
+    return automicrobatch_fsdp_hook_handles
+
+
 def _validate_evaluator(evaluator: Evaluator, device: Device):
     """Ensure automicrobatching is only on GPU.
 
@@ -897,7 +975,7 @@ class Trainer:
             (default: ``False``)
         autoresume (bool, optional): Whether or not to enable autoresume, which allows for stopping and resuming
             training. This allows use of spot instances, as the training run is now fault tolerant.  This parameter
-            requires ``save_folder`` and ``run_name`` to be specified and ``save_overwrite`` to be ``False``.
+            requires ``save_folder`` and ``run_name`` to be specified.
             (default: ``False``)
 
             When enabled, the save_folder is checked for checkpoints of the format ``"{save_folder}/{save_latest_filename}"``,
@@ -1106,6 +1184,17 @@ def __init__(
         # compile config for PyTorch 2.0 or higher
         compile_config: Optional[dict[str, Any]] = None,
     ):
+        if deepspeed_config is not None:
+            warnings.warn(
+                VersionedDeprecationWarning(
+                    'The use of DeepSpeed for training new models in Composer is deprecated. Composer is tightly integrated with PyTorch FSDP '
+                    +
+                    'which provides similar functionality. Please use the `parallelism_config` parameter instead. Please open '
+                    + 'a GitHub issue if you need help migrating from DeepSpeed to FSDP.',
+                    remove_version='0.28.0',
+                ),
+            )
+
         self.auto_log_hparams = auto_log_hparams
         self.python_log_level = python_log_level
         if self.python_log_level is not None:
@@ -1168,6 +1257,11 @@ def __init__(
                 '`Sequence parallelism requires a microbatch size of 1 distributed over the sequence parallel group.',
             )
 
+        # Automicrobatching
+        self.cumulative_alloc_retries = 0
+        self.num_consecutive_thrashes = 0
+        self.num_consecutive_non_OOM_batches = 0
+
         if auto_microbatching and profiler:
             raise ValueError(
                 "`device_train_microbatch_size='auto'` is not compatible with the profiler. It is "
@@ -1251,6 +1345,8 @@ def __init__(
         if parallelism_config is not None:
             # Patch PyTorch to fix distributed bugs
             patch_pytorch()
+            if auto_microbatching:
+                patch_unshard_for_automicrobatching(auto_microbatch_size_found=False)
 
         # Reproducibility
         rank_zero_seed, seed = _distribute_and_get_random_seed(seed, device)
@@ -1661,6 +1757,7 @@ def __init__(
             with reproducibility.seed_context(self.state.rank_zero_seed):
                 prepare_tp_module(
                     model,
+                    optimizers,
                     self.state.tp_config,
                 )
 
@@ -1668,7 +1765,7 @@ def __init__(
         if self.state.fsdp_config is not None and self.state.fsdp_config.auto_wrap and not self.state.load_monolith_rank0_only:
             # Init with globally fixed seed so all HSDP replicas have the same initial weights
             with reproducibility.seed_context(self.state.rank_zero_seed):
-                prepare_fsdp_module(
+                self.state.automicrobatch_fsdp_hook_handles, self.state.fsdp_modules = prepare_fsdp_module(
                     model,
                     optimizers,
                     self.state.fsdp_config,
@@ -1829,6 +1926,7 @@ def __init__(
                 algorithm_passes=self.engine.algorithm_passes,
             )
             self.state.run_name = run_name
+            self.state.load_path = load_path
 
         # FSDP wrap if model is not yet wrapped and FSDP is enabled. This can happen if
         # load_monolith_rank0_only=True but no checkpoint was loaded.
@@ -1838,7 +1936,14 @@ def __init__(
         ):
             # Init with globally fixed seed so all HSDP replicas have the same initial weights
             with reproducibility.seed_context(self.state.rank_zero_seed):
-                prepare_fsdp_module(model, optimizers, self.state.fsdp_config, precision, device, auto_microbatching)
+                self.state.automicrobatch_fsdp_hook_handles, self.state.fsdp_modules = prepare_fsdp_module(
+                    model,
+                    optimizers,
+                    self.state.fsdp_config,
+                    precision,
+                    device,
+                    auto_microbatching,
+                )
 
         self.engine.run_event(Event.AFTER_LOAD)
 
@@ -1986,7 +2091,7 @@ def _get_autoresume_checkpoint(
 
             signal_file_path = os.path.join(
                 os.path.dirname(latest_checkpoint_path),
-                f'.node_{dist.get_node_rank()}_local_rank0_completed_autoresume',
+                dist.get_node_signal_file_name(),
             )
             if dist.get_local_rank() == 0:
                 os.makedirs(os.path.dirname(signal_file_path), exist_ok=True)
@@ -2735,6 +2840,12 @@ def _train_batch(self, use_grad_scaling: bool) -> dict[str, torch.Tensor]:
         # Any in-place changes to a microbatch will be reflected in the device batch.
         device_batch = self.state.batch
 
+        # Define sync hook for FSDP modules if automicrobatching is on
+        sync_hook = _create_sync_hook(self.state)
+
+        original_microbatch_size = self.state.device_train_microbatch_size
+        oom_found_this_batch = False
+
         # Retry until we successfully complete training and return loss
         while True:
             # Reset train_metrics on every batch
@@ -2773,7 +2884,10 @@ def _train_batch(self, use_grad_scaling: bool) -> dict[str, torch.Tensor]:
                             else:
                                 optimizer.step()
             except RuntimeError as e:
-                if self.state.auto_microbatching and _is_cuda_oom(e):
+                if self.state.auto_microbatching and str(e) == OOM_FOUND_ON_OTHER_RANK:
+                    log.debug((f"A Different Rank OOM'd."))
+                    found_cuda_oom = 1
+                elif self.state.auto_microbatching and _is_cuda_oom(e):
                     log.debug((f"Rank {dist.get_global_rank()} OOM'd."))
                     found_cuda_oom = 1
                 elif self.state.auto_microbatching and ('cuda' in str(e).lower() or 'c10' in str(e).lower()):
@@ -2801,11 +2915,54 @@ def _train_batch(self, use_grad_scaling: bool) -> dict[str, torch.Tensor]:
                     dist.all_reduce(all_ranks_finished_tensor, reduce_operation='MIN')
                     all_ranks_finished = all_ranks_finished_tensor.item() == 1
                 if found_cuda_oom == 1:
+                    # Readd sync hooks if they were previously turned off
+                    if self.state.fsdp_enabled and len(self.state.automicrobatch_fsdp_hook_handles) == 0:
+                        self.state.automicrobatch_fsdp_hook_handles = _readd_fsdp_sync_hooks(
+                            self.state.fsdp_modules,
+                            sync_hook,
+                        )
                     _adjust_device_train_microbatch_size(self.state)
+                    self.num_consecutive_thrashes = 0
+                    self.num_consecutive_non_OOM_batches = 0
+                    oom_found_this_batch = True
                     # Skip return and rerun after handling oom
                     continue
+                if not oom_found_this_batch and torch.cuda.is_available():
+                    # Sync across all ranks to check if any rank had additional alloc retries this batch
+                    self.num_consecutive_thrashes = _update_num_consecutive_thrashes(
+                        self.state,
+                        self.num_consecutive_thrashes,
+                        self.cumulative_alloc_retries,
+                    )
+                if self.num_consecutive_thrashes >= 2:
+                    # Readd sync hooks if they were previously turned off
+                    if self.state.fsdp_enabled and len(self.state.automicrobatch_fsdp_hook_handles) == 0:
+                        self.state.automicrobatch_fsdp_hook_handles = _readd_fsdp_sync_hooks(
+                            self.state.fsdp_modules,
+                            sync_hook,
+                        )
+                    _adjust_device_train_microbatch_size(self.state)
+                    self.num_consecutive_thrashes = 0
+                    continue
+
             # Log microbatch and return loss if we've completed without OOMing.
             assert self.state.device_train_microbatch_size is not None
+            if original_microbatch_size != self.state.device_train_microbatch_size:
+                log.info(
+                    'Automicrobatching changed the microbatch size from '
+                    f'{original_microbatch_size} -> {self.state.device_train_microbatch_size}.',
+                )
+            self.num_consecutive_non_OOM_batches += 1
+            if self.state.fsdp_enabled and len(
+                self.state.automicrobatch_fsdp_hook_handles,
+            ) > 0 and self.num_consecutive_non_OOM_batches >= 3:
+                patch_unshard_for_automicrobatching(auto_microbatch_size_found=True)
+                for handle in self.state.automicrobatch_fsdp_hook_handles:
+                    handle.remove()
+                self.state.automicrobatch_fsdp_hook_handles.clear()
+            if torch.cuda.is_available():
+                memory_stats = torch.cuda.memory_stats()
+                self.cumulative_alloc_retries = memory_stats['num_alloc_retries']
             self.logger.log_metrics({'trainer/device_train_microbatch_size': self.state.device_train_microbatch_size})
             self.first_batch_complete = True
             return total_loss_dict
@@ -2976,7 +3133,7 @@ def _train_microbatch(
                 dist.all_reduce(all_ranks_finished_tensor, reduce_operation='MIN')
 
                 if found_cuda_oom == 1:
-                    raise RuntimeError('CUDA out of memory encountered on a different rank')
+                    raise RuntimeError(OOM_FOUND_ON_OTHER_RANK)
 
             # Loss
             self.engine.run_event(Event.BEFORE_LOSS)
@@ -3600,6 +3757,13 @@ def _eval_loop(
         if original_num_batches is not None:
             self.state.dataloader_len = original_num_batches
 
+        # If training occurs after evaluation, readd hooks in case of memory spike
+        if self.state.auto_microbatching:
+            sync_hook = _create_sync_hook(self.state)
+            if self.state.fsdp_enabled and len(self.state.automicrobatch_fsdp_hook_handles) == 0:
+                self.state.automicrobatch_fsdp_hook_handles = _readd_fsdp_sync_hooks(self.state.fsdp_modules, sync_hook)
+            self.num_consecutive_non_OOM_batches = 0
+
     def _use_grad_scaling(self, precision: Union[str, Precision], scaler: Optional[GradScaler]) -> bool:
         """Determines based on precision when to use grad scaling.
 
diff --git a/composer/utils/checkpoint.py b/composer/utils/checkpoint.py
index f9ad516724..56b13fcac6 100644
--- a/composer/utils/checkpoint.py
+++ b/composer/utils/checkpoint.py
@@ -36,6 +36,8 @@
     format_name_with_dist_and_time,
     get_file,
     is_tar,
+    is_uri,
+    maybe_create_object_store_from_uri,
     parse_uri,
 )
 from composer.utils.misc import ParallelismType, is_model_deepspeed, partial_format
@@ -214,9 +216,16 @@ def __init__(
         self,
         source_path: str,
         destination_path: str,
-        object_store: Union[ObjectStore, LoggerDestination],
-        device_mesh: Optional[DeviceMesh],
+        object_store: Optional[Union[ObjectStore, LoggerDestination]] = None,
+        device_mesh: Optional[DeviceMesh] = None,
     ):
+
+        if object_store is None:
+            if not is_uri(source_path):
+                raise ValueError('When object_store is None, source_path must be a URI.')
+            object_store = maybe_create_object_store_from_uri(source_path)
+            _, _, source_path = parse_uri(source_path)
+
         self.source_path = source_path
         self.destination_path = destination_path
         self.object_store = object_store
@@ -227,6 +236,7 @@ def __init__(
         metadata_destination = os.path.join(self.destination_path, '.metadata')
         if dist.get_local_rank() == 0:
             metadata_path = str(Path(source_path) / Path('.metadata'))
+            assert object_store is not None
             download_object_or_file(metadata_path, metadata_destination, object_store)
         dist.barrier()
 
@@ -274,6 +284,7 @@ def read_data(self, plan: LoadPlan, planner: LoadPlanner):
                     if not is_downloaded and not os.path.exists(file_destination):
                         log.debug(f'Downloading {relative_file_path} to {file_destination}.')
                         object_name = str(Path(self.source_path) / Path(relative_file_path))
+                        assert self.object_store is not None
                         download_object_or_file(object_name, file_destination, self.object_store)
                         log.debug(f'Finished downloading {relative_file_path} to {file_destination}.')
             except Exception as e:
@@ -819,7 +830,7 @@ def download_checkpoint(
         if not checkpoint_is_sharded:
             signal_file_path = os.path.join(
                 node_checkpoint_folder,
-                f'.node_{dist.get_node_rank()}_local_rank0_completed',
+                dist.get_node_signal_file_name(),
             )
             if dist.get_local_rank() == 0:
                 with open(signal_file_path, 'wb') as f:
diff --git a/composer/utils/dist.py b/composer/utils/dist.py
index 95a95835f4..2178ce2dd5 100644
--- a/composer/utils/dist.py
+++ b/composer/utils/dist.py
@@ -644,7 +644,7 @@ def get_node_signal_file_name(rng: Optional[random.Random] = None) -> str:
     random_string = ''.join(rng.choices(string.ascii_letters + string.digits, k=6))
     node_rank = get_node_rank()
     file_name_list = [f'._signal_file_node{node_rank}_{random_string}']
-    dist.broadcast_object_list(file_name_list, src=0)
+    broadcast_object_list(file_name_list, src=0)
     return file_name_list[0]
 
 
diff --git a/composer/utils/file_helpers.py b/composer/utils/file_helpers.py
index 4f458b0a8e..2d737a7fce 100644
--- a/composer/utils/file_helpers.py
+++ b/composer/utils/file_helpers.py
@@ -358,6 +358,19 @@ def format_name_with_dist_and_time(
 """
 
 
+def is_uri(path: str) -> bool:
+    """Check if the path is a URI.
+
+    Args:
+        path (str): The path to check.
+
+    Returns:
+        bool: Whether the path is a URI.
+    """
+    backend, _, _ = parse_uri(path)
+    return backend != ''
+
+
 def parse_uri(uri: str) -> tuple[str, str, str]:
     """Uses :py:func:`urllib.parse.urlparse` to parse the provided URI.
 
@@ -368,6 +381,8 @@ def parse_uri(uri: str) -> tuple[str, str, str]:
         tuple[str, str, str]: A tuple containing the backend (e.g. s3), bucket name, and path.
                               Backend name will be empty string if the input is a local path
     """
+    if isinstance(uri, pathlib.Path):
+        uri = str(uri)
     uri = uri.replace('AZURE_BLOBS', 'azure')  # urlparse does not support _ in scheme
     parse_result = urlparse(uri)
     backend, net_loc, path = parse_result.scheme, parse_result.netloc, parse_result.path
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 970af2f1ef..80ae8bad2e 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -172,7 +172,7 @@ RUN add-apt-repository ppa:deadsnakes/ppa && \
     rm -rf /var/lib/apt/lists/*
 
 RUN curl -fsSL https://bootstrap.pypa.io/get-pip.py | python${PYTHON_VERSION} - && \
-    pip${PYTHON_VERSION} install --no-cache-dir --upgrade 'pip<23' 'setuptools<70.0.0'
+    pip${PYTHON_VERSION} install --no-cache-dir --upgrade pip 'setuptools<70.0.0'
 
 #################
 # Install Pytorch
@@ -261,7 +261,7 @@ RUN if [ -n "$MOFED_VERSION" ] ; then \
 RUN if [ -n "$CUDA_VERSION" ] ; then \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir ninja==1.11.1 && \
         pip${PYTHON_VERSION} install --upgrade --no-cache-dir --force-reinstall packaging==22.0 && \
-        MAX_JOBS=1 pip${PYTHON_VERSION} install --no-cache-dir --no-build-isolation flash-attn==2.5.8; \
+        MAX_JOBS=1 pip${PYTHON_VERSION} install --no-cache-dir --no-build-isolation flash-attn==2.6.2; \
         cd .. ; \
     fi
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 533ce95b78..6b94be17b8 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -305,6 +305,11 @@ def get_algorithms_metadata() -> dict[str, dict[str, str]]:
 
 html_context = {'metadata': get_algorithms_metadata()}
 
+# Sphinx context injection deprecation patch
+html_baseurl = os.environ.get('READTHEDOCS_CANONICAL_URL', '')
+if os.environ.get('READTHEDOCS', '') == 'True':
+    html_context['readthedocs'] = True
+
 # ClassDocumenter.add_directive_header uses ClassDocumenter.add_line to
 #   write the class documentation.
 # We'll monkeypatch the add_line method and intercept lines that begin
diff --git a/pyproject.toml b/pyproject.toml
index 8ca97bc494..f153616f0d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -88,7 +88,7 @@ reportUnusedCoroutine = "error"
 # Pytest
 [tool.pytest.ini_options]
 # By default, do not run gpu, vision, docs, notebook, or daily tests
-addopts = "--codeblocks --strict-markers -m 'not gpu and not doctest and not daily and not remote'"
+addopts = "--codeblocks --strict-markers -m 'not gpu and not doctest and not daily and not remote' --color=yes"
 
 markers = [
     # Tests that require a world_size of two should be annotated with `@pytest.mark.world_size(2)`.
@@ -160,6 +160,10 @@ filterwarnings = [
     '''ignore:.*an autograd kernel was not registered to the Autograd key.*:UserWarning''',
     # Ignore save_state_dict / load_state_dict deprecation warnings
     '''ignore:'.*_state_dict' is deprecated and will be removed in future versions.*:UserWarning''',
+    # Ignore mlflow warnings about transformers versions,
+    '''ignore:The 'transformers' MLflow Models integration.*:UserWarning''',
+    # Ignore our own deprecation warnings,
+    '''ignore::composer.utils.warnings.VersionedDeprecationWarning''',
 ]
 
 # Coverage
diff --git a/setup.py b/setup.py
index c16807aa9d..2308781526 100644
--- a/setup.py
+++ b/setup.py
@@ -89,7 +89,7 @@ def package_files(prefix: str, directory: str, extension: str):
     'tabulate==0.9.0',  # for auto-generating tables
     'py-cpuinfo>=8.0.0,<10',
     'packaging>=21.3.0,<24.2',
-    'importlib-metadata>=5.0.0,<7',
+    'importlib-metadata>=5.0.0,<9',
     'mosaicml-cli>=0.5.25,<0.7',
     'pillow>=10.3.0,<11',
 ]
@@ -179,13 +179,13 @@ def package_files(prefix: str, directory: str, extension: str):
 ]
 
 extra_deps['nlp'] = [
-    'transformers>=4.11,!=4.34.0,<4.43',
+    'transformers>=4.11,!=4.34.0,<4.44',
     'datasets>=2.4,<3',
-    'huggingface-hub>=0.21.2,<0.24',
+    'huggingface-hub>=0.21.2,<0.25',
 ]
 
 extra_deps['peft'] = [
-    'peft>=0.10.0,<0.12',
+    'peft>=0.10.0,<0.13',
 ]
 
 extra_deps['sentencepiece'] = [
diff --git a/tests/checkpoint/helpers.py b/tests/checkpoint/helpers.py
index 4915c3a150..08ee3067ff 100644
--- a/tests/checkpoint/helpers.py
+++ b/tests/checkpoint/helpers.py
@@ -16,7 +16,9 @@
 from composer.callbacks import SpeedMonitor
 from composer.core import State
 from composer.devices import Device, DeviceCPU, DeviceGPU
+from composer.distributed import prepare_fsdp_module
 from composer.models import ComposerModel
+from composer.utils.parallelism import FSDPConfig
 from tests.common.models import EvenSimplerMLP, SimpleComposerMLP
 
 __all__ = [
@@ -37,12 +39,14 @@ def init_state(
     rank_zero_seed=10,
     run_name='test_run',
     take_step=False,
+    wrap_with_raw_fsdp=False,
 ) -> State:
     model, optimizer = init_model_and_optimizer(
         use_fsdp=use_fsdp,
         use_composer_model=True,
         take_step=take_step,
         device=device,
+        wrap_with_raw_fsdp=wrap_with_raw_fsdp,
     )
 
     test_dataset_sd = {'test': 0}
@@ -87,6 +91,7 @@ def init_model_and_optimizer(
     use_fsdp=False,
     tensor_type='sharded_tensor',
     device='cuda',
+    wrap_with_raw_fsdp=False,
 ) -> Tuple[Union[ComposerModel, torch.nn.Module], torch.optim.Optimizer]:
     model, loss_fn = init_model(
         use_composer_model,
@@ -95,6 +100,7 @@ def init_model_and_optimizer(
         use_fsdp=use_fsdp,
         tensor_type=tensor_type,
         device=device,
+        wrap_with_raw_fsdp=wrap_with_raw_fsdp,
     )
 
     optimizer = init_optimizer(
@@ -120,6 +126,7 @@ def init_model(
     tensor_type='sharded_tensor',
     sync_module_states=True,
     cpu_offload=False,
+    wrap_with_raw_fsdp=False,
 ) -> Tuple[Union[ComposerModel, torch.nn.Module], Any]:
     if use_composer_model:
         model = SimpleComposerMLP(num_features=num_features, num_classes=num_classes, device=device)
@@ -133,7 +140,6 @@ def init_model(
             use_orig_params=True,
             sync_module_states=sync_module_states,  # To enable easy comparison between rank 0 unsharded model and full state dict
             cpu_offload=CPUOffload(offload_params=True) if cpu_offload else None,
-            device_id=torch.device('cpu') if device == 'cpu' else None,
         )
 
         if tensor_type == 'dtensor':
@@ -141,10 +147,14 @@ def init_model(
             device_mesh = init_device_mesh('cuda', (2,))
             fsdp_kwargs['device_mesh'] = device_mesh
 
-        model = FSDP(
-            model,
-            **fsdp_kwargs,
-        )
+        if wrap_with_raw_fsdp:
+            model = FSDP(model, **fsdp_kwargs)
+        else:
+            prepare_fsdp_module(
+                model,
+                optimizers=None,
+                fsdp_config=FSDPConfig(**fsdp_kwargs),
+            )
 
     return model, loss_fn
 
diff --git a/tests/checkpoint/test_download.py b/tests/checkpoint/test_download.py
index 98c937bac4..d48845c983 100644
--- a/tests/checkpoint/test_download.py
+++ b/tests/checkpoint/test_download.py
@@ -8,7 +8,7 @@
 import pytest
 import torch
 
-from composer.checkpoint import download_monolithic_checkpoint
+from composer.checkpoint import download_monolithic_checkpoint, get_model_state_dict
 from composer.utils import dist
 from tests.checkpoint.helpers import init_model
 from tests.common.markers import world_size
@@ -26,8 +26,7 @@ def test_download_monolithic_checkpoint(world_size: int, rank_zero_only: bool):
         use_fsdp = True
     fsdp_model, _ = init_model(use_fsdp=use_fsdp)
 
-    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
-    state = get_model_state_dict(fsdp_model, options=StateDictOptions(full_state_dict=True))
+    state = get_model_state_dict(fsdp_model, sharded_state_dict=False)
 
     checkpoint_filename = 'state_dict'
     save_filename = os.path.join(tmp_dir.name, checkpoint_filename)
diff --git a/tests/checkpoint/test_load.py b/tests/checkpoint/test_load.py
new file mode 100644
index 0000000000..aa4f50059e
--- /dev/null
+++ b/tests/checkpoint/test_load.py
@@ -0,0 +1,433 @@
+# Copyright 2024 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import os
+import uuid
+from pathlib import Path
+
+import pytest
+from torch.utils.data import DataLoader
+
+from composer.checkpoint.load import (
+    load_checkpoint,
+    load_model_checkpoint,
+    load_optim_checkpoint,
+    load_resumption_checkpoint,
+)
+from composer.checkpoint.save import (
+    save_checkpoint_to_disk,
+    save_model_to_disk,
+    save_optim_to_disk,
+    save_resumption_state_to_disk,
+)
+from composer.checkpoint.state_dict import (
+    _is_model_fsdp,
+    get_model_state_dict,
+    get_optim_state_dict,
+    get_resumption_state_dict,
+)
+from composer.trainer import Trainer
+from composer.utils import dist
+from tests.checkpoint.helpers import init_model, init_model_and_optimizer, init_state
+from tests.common import (
+    RandomClassificationDataset,
+)
+from tests.common.compare import deep_compare
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    'world_size,sharded_model,sharded_checkpoint,shard_as_needed_during_load',
+    [
+        #Loading an unsharded checkpoint into an unsharded model on a single GPU (not sharding after)
+        pytest.param(1, False, False, False, marks=pytest.mark.world_size(1)),
+
+        # Loading a sharded checkpoint into a sharded model in distributed setting
+        pytest.param(2, True, True, False, marks=pytest.mark.world_size(2)),
+
+        # SHOULD FAIL: Loading an unsharded checkpoint into a sharded model
+        pytest.param(2, True, False, False, marks=pytest.mark.world_size(2)),
+
+        # SHOULD FAIL: Attempting to load a sharded checkpoint into an unsharded model without sharding
+        pytest.param(2, False, True, False, marks=pytest.mark.world_size(2)),
+
+        # Loading a sharded checkpoint into an unsharded model (sharding it before load)
+        pytest.param(2, False, True, True, marks=pytest.mark.world_size(2)),
+
+        # Loading an unsharded checkpoint into an unsharded model and sharding it after.
+        pytest.param(2, False, False, True, marks=pytest.mark.world_size(2)),
+
+        # The other three permutations of the above tests are:
+        # 2 gpu, Sharded model, sharded checkpoint, with additional sharding -> no need to shard already sharded model
+        # 2 gpu, Sharded model, unsharded checkpoint, with additional sharding -> no need to shard already sharded model
+        # 2 gpu, Unsharded model, unsharded checkpoint, without additional sharding -> no need to try this on 2 gpus
+    ],
+)
+def test_load_model_checkpoint(
+    world_size: int,
+    tmp_path: Path,
+    sharded_model: bool,
+    sharded_checkpoint: bool,
+    shard_as_needed_during_load: bool,
+):
+    if sharded_model and not sharded_checkpoint:
+        pytest.xfail(
+            'Loading an unsharded checkpoint into a sharded model is not supported and causes OOMs when running with these tests',
+        )
+    # Ensure all ranks use the same path
+    destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
+    destination_dir = dist.all_gather_object(destination_dir)[0]
+
+    # Save a model checkpoint
+    model, _ = init_model(use_fsdp=sharded_checkpoint, device='cuda')
+    save_path = os.path.join(destination_dir, 'model.pt') if not sharded_checkpoint else destination_dir
+    saved_path = save_model_to_disk(model, save_path, sharded_checkpoint=sharded_checkpoint)
+
+    # Get the original model's state dict
+    original_state_dict = get_model_state_dict(model, sharded_state_dict=False)
+    # Load the model checkpoint
+    new_model, _ = init_model(use_fsdp=sharded_model, device='cuda')
+    if saved_path is not None:
+        load_path = saved_path if not sharded_checkpoint else str(Path(saved_path).parent)
+    else:
+        load_path = ''
+
+    if not sharded_model and sharded_checkpoint and not shard_as_needed_during_load:
+        context_manager = pytest.raises(ValueError)
+    else:
+        context_manager = contextlib.nullcontext()
+
+    with context_manager:
+        load_model_checkpoint(
+            new_model,
+            load_path=load_path,
+            load_options=dict(
+                sharded_checkpoint=sharded_checkpoint,
+                shard_as_needed_during_load=shard_as_needed_during_load,
+            ),
+        )
+        # Check if model is sharded when it should be
+        if shard_as_needed_during_load:
+            assert _is_model_fsdp(new_model), 'Model should be sharded after load'
+
+        # Get the new model's state dict
+        new_state_dict = get_model_state_dict(new_model, sharded_state_dict=False)
+
+        if dist.get_global_rank() == 0:
+            deep_compare(original_state_dict, new_state_dict)
+
+
+@pytest.mark.filterwarnings('ignore:TypedStorage is deprecated.')
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    'world_size,sharded_optimizer,sharded_checkpoint,shard_as_needed_during_load',
+    [
+        # Loading an unsharded checkpoint into an unsharded optimizer on a single GPU (not sharding after)
+        pytest.param(1, False, False, False, marks=pytest.mark.world_size(1)),
+
+        # Loading a sharded checkpoint into a sharded optimizer in distributed setting
+        pytest.param(2, True, True, False, marks=pytest.mark.world_size(2)),
+    ],
+)
+def test_load_optim_checkpoint(
+    world_size: int,
+    tmp_path: Path,
+    sharded_optimizer: bool,
+    sharded_checkpoint: bool,
+    shard_as_needed_during_load: bool,
+):
+    if sharded_optimizer and not sharded_checkpoint:
+        pytest.xfail(
+            'Loading an unsharded checkpoint into a sharded optimizer is not supported and causes OOMs when running with these tests',
+        )
+
+    # Ensure all ranks use the same path
+    destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
+    destination_dir = dist.all_gather_object(destination_dir)[0]
+
+    # Save an optimizer checkpoint
+    model, optimizer = init_model_and_optimizer(use_fsdp=sharded_checkpoint, device='cuda')
+    save_path = os.path.join(destination_dir, 'optim.pt') if not sharded_checkpoint else destination_dir
+    saved_path = save_optim_to_disk(model, optimizer, save_path, sharded_checkpoint=sharded_checkpoint)
+
+    # Get the original optimizer's state dict
+    original_state_dict = get_optim_state_dict(model, optimizer, sharded_state_dict=False)
+
+    # Load the optimizer checkpoint
+    new_model, new_optimizer = init_model_and_optimizer(use_fsdp=sharded_optimizer, device='cuda')
+    if saved_path is not None:
+        load_path = saved_path if not sharded_checkpoint else str(Path(saved_path).parent)
+    else:
+        load_path = ''
+
+    if not sharded_optimizer and sharded_checkpoint and not shard_as_needed_during_load:
+        context_manager = pytest.raises(ValueError)
+    else:
+        context_manager = contextlib.nullcontext()
+
+    with context_manager:
+        load_optim_checkpoint(
+            new_model,
+            new_optimizer,
+            load_path=load_path,
+            load_options=dict(
+                sharded_checkpoint=sharded_checkpoint,
+                shard_as_needed_during_load=shard_as_needed_during_load,
+            ),
+        )
+
+        # Check if optimizer is sharded when it should be
+        if shard_as_needed_during_load:
+            assert _is_model_fsdp(new_model), 'Optimizer should be sharded after load'
+
+        # Get the new optimizer's state dict
+        new_state_dict = get_optim_state_dict(new_model, new_optimizer, sharded_state_dict=False)
+
+        if dist.get_global_rank() == 0:
+            deep_compare(original_state_dict, new_state_dict)
+
+
+@pytest.mark.gpu
+@pytest.mark.filterwarnings('ignore:SWA has')
+def test_load_resumption_checkpoint(tmp_path: Path):
+    # Ensure all ranks use the same path
+    destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
+    destination_dir = dist.all_gather_object(destination_dir)[0]
+
+    # Create an initial state using the helper function
+    initial_state = init_state(
+        device='cpu',  # or 'cuda' if you want to test on GPU
+        include_schedulers=True,
+        include_algorithms=True,
+        include_callbacks=True,
+        use_grad_scaler=True,
+        rank_zero_seed=42,
+        run_name='test_run',
+    )
+
+    # Modify some values to ensure they're changed
+    initial_state.timestamp.to_next_batch()
+    initial_state.dataset_state = {'train': {'some_key': 'some_value'}}
+
+    # Save the resumption state
+    save_path = os.path.join(destination_dir, 'resumption.pkl')
+    save_resumption_state_to_disk(initial_state, save_path)
+
+    # Create a new state
+    new_state = init_state(
+        device='cpu',  # or 'cuda' if you want to test on GPU
+        include_schedulers=True,
+        include_algorithms=True,
+        include_callbacks=True,
+        use_grad_scaler=True,
+        rank_zero_seed=42,
+        run_name='test_run',
+    )
+    import time
+    time.sleep(1)
+    # Load the resumption checkpoint
+    load_resumption_checkpoint(new_state, save_path)
+
+    # Check that the loaded state matches the initial state
+    deep_compare(new_state.timestamp.state_dict(), initial_state.timestamp.state_dict())
+    assert new_state.rank_zero_seed == initial_state.rank_zero_seed
+    assert new_state.run_name == initial_state.run_name
+
+    # Check schedulers, algorithms, callbacks, and scaler
+    if initial_state.schedulers:
+        for init_scheduler, new_scheduler in zip(initial_state.schedulers, new_state.schedulers):
+            deep_compare(init_scheduler.state_dict(), new_scheduler.state_dict())
+
+    if initial_state.algorithms:
+        for init_algo, new_algo in zip(initial_state.algorithms, new_state.algorithms):
+            init_state_dict = init_algo.state_dict()
+            new_state_dict = new_algo.state_dict()
+            if 'repr' in init_state_dict:
+                init_state_dict.pop('repr')
+            if 'repr' in new_state_dict:
+                new_state_dict.pop('repr')
+            deep_compare(init_state_dict, new_state_dict)
+
+    if initial_state.callbacks:
+        for init_callback, new_callback in zip(initial_state.callbacks, new_state.callbacks):
+            deep_compare(init_callback.state_dict(), new_callback.state_dict())
+
+    if initial_state.scaler:
+        assert initial_state.scaler is not None
+        assert new_state.scaler is not None
+        deep_compare(initial_state.scaler.state_dict(), new_state.scaler.state_dict())
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    'world_size,sharded_model,sharded_checkpoint,shard_as_needed_during_load',
+    [
+        # Loading an unsharded checkpoint into an unsharded model on a single GPU (not sharding after)
+        pytest.param(1, False, False, False, marks=pytest.mark.world_size(1)),
+
+        # Loading a sharded checkpoint into a sharded model in distributed setting
+        pytest.param(2, True, True, False, marks=pytest.mark.world_size(2)),
+
+        # SHOULD FAIL: Loading an unsharded checkpoint into a sharded model
+        pytest.param(2, True, False, False, marks=pytest.mark.world_size(2)),
+
+        # SHOULD FAIL: Attempting to load a sharded checkpoint into an unsharded model without sharding
+        pytest.param(2, False, True, False, marks=pytest.mark.world_size(2)),
+    ],
+)
+def test_load_checkpoint(
+    world_size: int,
+    tmp_path: Path,
+    sharded_model: bool,
+    sharded_checkpoint: bool,
+    shard_as_needed_during_load: bool,
+):
+
+    if sharded_model and not sharded_checkpoint:
+        pytest.xfail(
+            'Loading an unsharded checkpoint into a sharded model is not supported and causes OOMs when running with these tests',
+        )
+    # Ensure all ranks use the same path
+    destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
+    destination_dir = dist.all_gather_object(destination_dir)[0]
+
+    # Save an optimizer checkpoint
+    state = init_state(
+        use_fsdp=sharded_checkpoint,
+        device='cuda',
+        take_step=True,
+    )
+    load_path = save_checkpoint_to_disk(
+        destination_dir=destination_dir,
+        state=state,
+        options={
+            'sharded_checkpoint': sharded_checkpoint,
+            'save_model': True,
+            'save_optimizer': True,
+            'save_resumption_state': True,
+        },
+    )
+    original_model_state_dict = get_model_state_dict(state.model, sharded_state_dict=False)
+    original_optim_state_dict = get_optim_state_dict(state.model, state.optimizers[0], sharded_state_dict=False)
+    original_resumption_state = get_resumption_state_dict(state)
+    new_state = init_state(use_fsdp=sharded_model, device='cuda', take_step=True)
+    if not sharded_model and sharded_checkpoint and not shard_as_needed_during_load:
+        context_manager = pytest.raises(ValueError)
+    else:
+        context_manager = contextlib.nullcontext()
+    with context_manager:
+        load_checkpoint(
+            load_path=load_path,
+            state=new_state,
+            load_options={
+                'sharded_checkpoint': sharded_checkpoint,
+                'load_optimizer': True,
+                'load_resumption_state': True,
+                'shard_as_needed_during_load': shard_as_needed_during_load,
+            },
+        )
+        if shard_as_needed_during_load:
+            assert _is_model_fsdp(new_state.model), 'Model should be sharded after load'
+        # Get the new model's state dict
+        new_state_dict = get_model_state_dict(new_state.model, sharded_state_dict=False)
+        new_optim_state_dict = get_optim_state_dict(new_state.model, new_state.optimizers[0], sharded_state_dict=False)
+        new_resumption_state = get_resumption_state_dict(new_state)
+
+        if dist.get_global_rank() == 0:
+            deep_compare(original_model_state_dict, new_state_dict)
+            deep_compare(original_optim_state_dict, new_optim_state_dict)
+            deep_compare(original_resumption_state, new_resumption_state, ignore_keys=['rng', 'run_name'])
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    'world_size,sharded_model,sharded_checkpoint,shard_as_needed_during_load',
+    [
+        # Loading an unsharded checkpoint into an unsharded model on a single GPU (not sharding after)
+        pytest.param(1, False, False, False, marks=pytest.mark.world_size(1)),
+
+        # Loading a sharded checkpoint into a sharded model in distributed setting
+        pytest.param(2, True, True, False, marks=pytest.mark.world_size(2)),
+
+        # Loading a sharded checkpoint into an unsharded model (sharding it before load)
+        pytest.param(2, False, True, True, marks=pytest.mark.world_size(2)),
+
+        # Loading an unsharded checkpoint into an unsharded model and sharding it after.
+        pytest.param(2, False, False, True, marks=pytest.mark.world_size(2)),
+
+        # The other three permutations of the above tests are:
+        # 2 gpu, Sharded model, sharded checkpoint, with additional sharding -> no need to shard already sharded model
+        # 2 gpu, Sharded model, unsharded checkpoint, with additional sharding -> no need to shard already sharded model
+        # 2 gpu, Unsharded model, unsharded checkpoint, without additional sharding -> no need to try this on 2 gpus
+    ],
+)
+def test_load_model_checkpoint_and_eval(
+    world_size: int,
+    tmp_path: Path,
+    sharded_model: bool,
+    sharded_checkpoint: bool,
+    shard_as_needed_during_load: bool,
+):
+    if sharded_model and not sharded_checkpoint:
+        pytest.xfail(
+            'Loading an unsharded checkpoint into a sharded model is not supported and causes OOMs when running with these tests',
+        )
+    # Ensure all ranks use the same path
+    destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
+    destination_dir = dist.all_gather_object(destination_dir)[0]
+
+    # Save a model checkpoint
+    model, _ = init_model(use_composer_model=True, use_fsdp=sharded_checkpoint, device='cuda')
+    save_path = os.path.join(destination_dir, 'model.pt') if not sharded_checkpoint else destination_dir
+    saved_path = save_model_to_disk(model, save_path, sharded_checkpoint=sharded_checkpoint)
+
+    # Get the original model's state dict
+    original_state_dict = get_model_state_dict(model, sharded_state_dict=False)
+    # Load the model checkpoint
+    new_model, _ = init_model(use_composer_model=True, use_fsdp=sharded_model, device='cuda')
+    if saved_path is not None:
+        load_path = saved_path if not sharded_checkpoint else str(Path(saved_path).parent)
+    else:
+        load_path = ''
+
+    if not sharded_model and sharded_checkpoint and not shard_as_needed_during_load:
+        context_manager = pytest.raises(ValueError)
+    else:
+        context_manager = contextlib.nullcontext()
+
+    with context_manager:
+        load_model_checkpoint(
+            new_model,
+            load_path=load_path,
+            load_options=dict(
+                sharded_checkpoint=sharded_checkpoint,
+                shard_as_needed_during_load=shard_as_needed_during_load,
+            ),
+        )
+        # Check if model is sharded when it should be
+        if shard_as_needed_during_load:
+            assert _is_model_fsdp(new_model), 'Model should be sharded after load'
+
+        # Get the new model's state dict
+        new_state_dict = get_model_state_dict(new_model, sharded_state_dict=False)
+
+        if dist.get_global_rank() == 0:
+            deep_compare(original_state_dict, new_state_dict)
+
+        dataset = RandomClassificationDataset(
+            shape=(8,),
+            size=100,
+            num_classes=3,
+        )
+
+        trainer = Trainer(
+            eval_dataloader=DataLoader(
+                dataset=dataset,
+                sampler=dist.get_sampler(dataset),
+            ),
+            model=new_model,  # type: ignore
+        )
+
+        # Evaluate the model
+        trainer.eval()
diff --git a/tests/checkpoint/test_save.py b/tests/checkpoint/test_save.py
index f4d41cc09d..04a3122b06 100644
--- a/tests/checkpoint/test_save.py
+++ b/tests/checkpoint/test_save.py
@@ -46,7 +46,6 @@ def test_save_checkpoint_to_disk(world_size: int, tmp_path: str, sharded_model:
     destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
     destination_dir = dist.all_gather_object(destination_dir)[0]
     save_options = {
-        'destination_dir': destination_dir,
         'save_model': True,
         'save_optimizer': True,
         'save_resumption_state': True,
@@ -57,7 +56,7 @@ def test_save_checkpoint_to_disk(world_size: int, tmp_path: str, sharded_model:
     state.run_name = 'foo'
     state.timestamp = Timestamp()
     expected_destination_dir = os.path.join(destination_dir, 'ep0-ba0')
-    save_checkpoint_to_disk(state, save_options)
+    save_checkpoint_to_disk(destination_dir=destination_dir, state=state, options=save_options)
     expected_model_dir = os.path.join(expected_destination_dir, 'model')
     expected_optim_dir = os.path.join(expected_destination_dir, 'optim')
     expected_metadata_filepath = os.path.join(expected_destination_dir, 'composer_metadata.json')
@@ -84,8 +83,9 @@ def test_save_checkpoint_to_disk(world_size: int, tmp_path: str, sharded_model:
 def test_save_composer_metadata_to_disk(tmp_path: str):
     destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
     destination_dir = dist.all_gather_object(destination_dir)[0]
+    destination_dir = os.path.join(destination_dir, 'composer_metadata.json')
     save_composer_metadata_to_disk(destination_dir)
-    expected_file_path = os.path.join(destination_dir, 'composer_metadata.json')
+    expected_file_path = destination_dir
     assert os.path.exists(expected_file_path)
     json.load(open(expected_file_path, 'r'))
 
@@ -103,6 +103,7 @@ def test_save_optim_to_disk(world_size: int, tmp_path: str, sharded_optimizer: b
     destination_dir = os.path.join(tmp_path, str(uuid.uuid4())[:8])
     # Sync the path across all ranks
     destination_dir = dist.all_gather_object(destination_dir)[0]
+    destination_dir = os.path.join(destination_dir, 'optim.pt') if not sharded_checkpoint else destination_dir
     use_fsdp = sharded_optimizer
     model, optim = init_model_and_optimizer(use_fsdp=use_fsdp, device='cuda')
     optim_state_dict = get_optim_state_dict(model, optimizer=optim, sharded_state_dict=sharded_checkpoint)
@@ -114,14 +115,14 @@ def test_save_optim_to_disk(world_size: int, tmp_path: str, sharded_optimizer: b
     cur_state_dict = get_optim_state_dict(model, optimizer=optim, sharded_state_dict=sharded_checkpoint)
 
     if sharded_checkpoint:
-        expected_file_path = os.path.join(destination_dir, 'optim')
+        expected_file_path = os.path.join(destination_dir)
         if version.parse(torch.__version__) < version.parse('2.2.0'):
             DCP.load_state_dict(state_dict=cur_state_dict, storage_reader=DCP.FileSystemReader(expected_file_path))
         else:
             DCP.load(state_dict=cur_state_dict, storage_reader=DCP.FileSystemReader(expected_file_path))
     else:
         if dist.get_global_rank() == 0:
-            expected_file_path = os.path.join(destination_dir, 'optim', 'optim.pt')
+            expected_file_path = destination_dir
             cur_state_dict = torch.load(expected_file_path, map_location='cuda')
 
     deep_compare(optim_state_dict_saved, cur_state_dict)
@@ -144,6 +145,7 @@ def test_save_model_to_disk(world_size: int, tmp_path: str, sharded_model: bool,
     model, _ = init_model(use_fsdp=use_fsdp, device='cuda', sync_module_states=True)
     state_dict = get_model_state_dict(model, sharded_state_dict=sharded_checkpoint)
     state_dict_saved = deepcopy(state_dict)
+    destination_dir = os.path.join(destination_dir, 'model.pt') if not sharded_checkpoint else destination_dir
     save_model_to_disk(model, destination_dir=destination_dir, sharded_checkpoint=sharded_checkpoint)
 
     # Load new model from disk
@@ -151,14 +153,14 @@ def test_save_model_to_disk(world_size: int, tmp_path: str, sharded_model: bool,
     cur_state_dict = get_model_state_dict(new_model, sharded_state_dict=sharded_checkpoint)
 
     if sharded_checkpoint:
-        expected_file_path = os.path.join(destination_dir, 'model')
+        expected_file_path = destination_dir
         if version.parse(torch.__version__) < version.parse('2.2.0'):
             DCP.load_state_dict(state_dict=cur_state_dict, storage_reader=DCP.FileSystemReader(expected_file_path))
         else:
             DCP.load(state_dict=cur_state_dict, storage_reader=DCP.FileSystemReader(expected_file_path))
     else:
         if dist.get_global_rank() == 0:
-            expected_file_path = os.path.join(destination_dir, 'model', 'model.pt')
+            expected_file_path = destination_dir
             cur_state_dict = torch.load(expected_file_path, map_location='cuda')
 
     deep_compare(state_dict_saved, cur_state_dict)
diff --git a/tests/checkpoint/test_state_dict.py b/tests/checkpoint/test_state_dict.py
index 12fde27249..5a316ee286 100644
--- a/tests/checkpoint/test_state_dict.py
+++ b/tests/checkpoint/test_state_dict.py
@@ -243,7 +243,11 @@ def test_get_model_state_dict_precision_unsharded_model(precision: str, use_comp
 @pytest.mark.gpu
 @pytest.mark.parametrize('use_composer_model', [True, False])
 def test_get_optim_state_dict_unsharded_model(use_composer_model: bool):
-    model, optimizer = init_model_and_optimizer(use_composer_model=use_composer_model, take_step=True)
+    model, optimizer = init_model_and_optimizer(
+        use_composer_model=use_composer_model,
+        take_step=True,
+        wrap_with_raw_fsdp=True,
+    )
     optim_state_dict = get_optim_state_dict(model, optimizer)
 
     # Dict mapping parameter index to optimizer state for that parameter.
@@ -287,7 +291,11 @@ def test_get_optim_state_dict_unsharded_model(use_composer_model: bool):
 )
 @pytest.mark.parametrize('use_composer_model', [True, False])
 def test_get_optim_state_dict_precision_unsharded_model(precision: str, use_composer_model: bool):
-    model, optimizer = init_model_and_optimizer(use_composer_model=use_composer_model, take_step=True)
+    model, optimizer = init_model_and_optimizer(
+        use_composer_model=use_composer_model,
+        take_step=True,
+        wrap_with_raw_fsdp=True,
+    )
     optim_state_dict = get_optim_state_dict(model, optimizer, precision=precision)
     for param_state in optim_state_dict['state'].values():
         assert param_state['exp_avg'].dtype == precision
@@ -307,6 +315,7 @@ def test_get_optim_dict_full_for_sharded_model(world_size, tensor_type, use_comp
         take_step=True,
         use_fsdp=True,
         tensor_type=tensor_type,
+        wrap_with_raw_fsdp=True,
     )
     optim_state_dict = get_optim_state_dict(model, optimizer, sharded_state_dict=False)
 
@@ -316,7 +325,10 @@ def test_get_optim_dict_full_for_sharded_model(world_size, tensor_type, use_comp
         if dist.get_global_rank() == 0:
             # Because model is sharded, the state dict should have the same keys as the model's parameters.
             for fqn, param_state in optim_state_dict['state'].items():
-                model_param_shape = fqn_to_shape_map[fqn]
+                try:
+                    model_param_shape = fqn_to_shape_map[fqn]
+                except KeyError:
+                    assert False, f'{fqn_to_shape_map}'
                 assert model_param_shape == param_state['exp_avg'].shape
                 assert model_param_shape == param_state['exp_avg_sq'].shape
 
@@ -334,6 +346,7 @@ def test_get_optim_dict_sharded_for_sharded_model(world_size, tensor_type, use_c
         take_step=True,
         use_fsdp=True,
         tensor_type=tensor_type,
+        wrap_with_raw_fsdp=True,
     )
     model_state_dict = get_model_state_dict(model, sharded_state_dict=True)
     optim_state_dict = get_optim_state_dict(model, optimizer, sharded_state_dict=True)
@@ -447,6 +460,7 @@ def test_get_resumption_state_dict():
         include_schedulers=True,
         rank_zero_seed=rank_zero_seed,
         run_name=run_name,
+        wrap_with_raw_fsdp=True,
     )
     test_dataset_sd = {'test': 0}
     rsd = get_resumption_state_dict(state)
@@ -489,11 +503,11 @@ def test_get_resumption_state_dict():
 
 @pytest.mark.gpu
 def test_get_resumption_state_dict_gpu():
-    state = init_state(device='cuda', use_grad_scaler=True)
+    state = init_state(device='cuda', use_grad_scaler=True, wrap_with_raw_fsdp=True)
     rsd = get_resumption_state_dict(state)
     assert 'scaler' in rsd
     assert set(
-        rsd['scaler'].keys(),
+        rsd['scaler']['GradScaler'].keys(),
     ) == {'scale', 'growth_factor', 'backoff_factor', 'growth_interval', '_growth_tracker'}
 
     assert 'rng' in rsd
diff --git a/tests/checkpoint/test_upload.py b/tests/checkpoint/test_upload.py
new file mode 100644
index 0000000000..6b4e8dda23
--- /dev/null
+++ b/tests/checkpoint/test_upload.py
@@ -0,0 +1,74 @@
+# Copyright 2024 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+import multiprocessing
+import os
+import tempfile
+from unittest.mock import patch
+
+import pytest
+
+from composer.checkpoint import upload_file
+from composer.core import Engine, Event, State
+from composer.loggers import Logger
+from composer.utils import dist
+from tests.common.markers import world_size
+from tests.utils.test_remote_uploader import DummyObjectStore
+
+
+@world_size(1, 2)
+@pytest.mark.gpu
+@pytest.mark.parametrize(('async_upload'), [True, False])
+def test_upload_file(world_size: int, async_upload, dummy_state: State):
+    # Prepare local file to upload
+    local_file_path = tempfile.TemporaryDirectory()
+    rank = dist.get_global_rank() if world_size > 1 else 0
+    file_name = f'checkpoint_file_rank_{rank}'
+    local_file_full_name = os.path.join(local_file_path.name, file_name)
+    with open(local_file_full_name, 'w') as f:
+        f.write(str(rank))
+
+    # Prepare remote path
+    remote_path = tempfile.TemporaryDirectory()
+    if world_size > 1:
+        remote_path_list = [remote_path]
+        dist.broadcast_object_list(remote_path_list, src=0)
+        remote_path = remote_path_list[0]
+
+    def _get_tmp_dir(self):
+        return remote_path
+
+    fork_context = multiprocessing.get_context('fork')
+    with patch('composer.utils.object_store.utils.S3ObjectStore', DummyObjectStore):
+        with patch('tests.utils.test_remote_uploader.DummyObjectStore.get_tmp_dir', _get_tmp_dir):
+            with patch('composer.utils.remote_uploader.multiprocessing.get_context', lambda _: fork_context):
+                symlink_file_name = 'latest.symlink'
+                engine = None
+                if async_upload:
+                    logger = Logger(dummy_state)
+                    engine = Engine(state=dummy_state, logger=logger)
+                path = 'save_path'
+                upload_file(
+                    source_path=local_file_full_name,
+                    dest_dir='S3://bucket_name/' + path,
+                    symlink_granularity='file',
+                    async_upload=async_upload,
+                    state=dummy_state,
+                )
+                if async_upload:
+                    assert engine is not None
+                    engine.run_event(Event.FIT_END)
+                if world_size > 1:
+                    dist.barrier()
+
+                remote_file_name = os.path.join(remote_path.name, path, file_name)
+                assert os.path.isfile(remote_file_name)
+                with open(remote_file_name, 'r') as f:
+                    assert f.read() == str(rank)
+                remote_symlink_file_name = os.path.join(remote_path.name, path, symlink_file_name)
+                if rank == 0:
+                    assert os.path.isfile(remote_symlink_file_name)
+                    with open(remote_symlink_file_name, 'r') as f:
+                        assert f.read() == path
+                if world_size > 1:
+                    dist.barrier()
diff --git a/tests/trainer/test_fsdp.py b/tests/trainer/test_fsdp.py
index b077d22131..7b9bd4825c 100644
--- a/tests/trainer/test_fsdp.py
+++ b/tests/trainer/test_fsdp.py
@@ -1,13 +1,14 @@
 # Copyright 2022 MosaicML Composer authors
 # SPDX-License-Identifier: Apache-2.0
 
-from unittest.mock import MagicMock
+import copy
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 from packaging import version
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 
 from composer.models import ComposerClassifier, ComposerModel
 from composer.trainer.trainer import Trainer, _fsdp_reshard_and_cleanup
@@ -206,6 +207,131 @@ def test_fsdp_prefetch_limit(forward_prefetch_limit: int, backward_prefetch_limi
     trainer.fit()
 
 
+class SimpleDatasetForAuto(Dataset):
+
+    def __init__(self, size: int = 256, feature_size: int = 1, num_classes: int = 2):
+        self.size = size
+        self.feature_size = feature_size
+        self.num_classes = num_classes
+        self.x = None
+        self.y = None
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, index: int):
+        # Note: lazily generate data so it runs after Composer seeds everything, giving the same
+        # dataset across multiple calls when using the same seed.
+        if self.x is None:
+            self.x = torch.randn(self.size, self.feature_size)
+        if self.y is None:
+            self.y = torch.randint(0, self.num_classes, size=(self.size,), dtype=torch.long)
+        return self.x[index]
+
+
+class SimpleMLPForTestingOOM(ComposerModel):
+
+    def __init__(self, num_features: int = 128, device: str = 'cuda'):
+        super().__init__()
+        self.device = device
+        self.fc1 = torch.nn.Linear(num_features, num_features, device=device, bias=False)
+        self.fc2 = torch.nn.Linear(num_features, num_features, device=device, bias=False)
+        self.fc3 = torch.nn.Linear(num_features, num_features, device=device, bias=False)
+        self.rank = dist.get_global_rank()
+        self.iter = 0
+
+    def forward(self, x):
+        x = self.fc1(x)
+        if self.rank == 0 and x.shape[0] >= 64:
+            raise RuntimeError('CUDA out of memory')
+        x = self.fc2(x)
+        x = self.fc3(x)
+        self.iter += 1
+        return x
+
+    def loss(self, outputs, batch):
+        return torch.sum(outputs)
+
+
+@pytest.mark.gpu
+@pytest.mark.filterwarnings("ignore:`device_train_microbatch_size='auto'` may potentially fail with unexpected.*")
+@pytest.mark.filterwarnings('ignore:CUDA out of memory*')
+@world_size(2)
+def test_automicrobatching_fsdp(world_size: int):
+    model = SimpleMLPForTestingOOM()
+    model.fc1._fsdp_wrap = True  # pyright: ignore[reportGeneralTypeIssues]
+    model.fc2._fsdp_wrap = True  # pyright: ignore[reportGeneralTypeIssues]
+    dataset = SimpleDatasetForAuto(size=256, feature_size=128)
+    train_dataloader = DataLoader(dataset, batch_size=64, sampler=dist.get_sampler(dataset))
+    trainer = Trainer(
+        model=model,
+        train_dataloader=train_dataloader,
+        fsdp_config={
+            'forward_prefetch_limit': 1,
+            'backward_prefetch_limit': 1,
+        },
+        max_duration='1ba',
+        device='gpu',
+        device_train_microbatch_size='auto',
+        dist_timeout=20,
+    )
+    trainer.fit()
+
+
+class SimpleMLPForTestingHooks(ComposerModel):
+
+    def __init__(self, num_features: int = 128, device: str = 'cuda'):
+        super().__init__()
+        self.device = device
+        self.fc1 = torch.nn.Linear(num_features, num_features, device=device, bias=False)
+        self.fc2 = torch.nn.Linear(num_features, num_features, device=device, bias=False)
+        self.fc3 = torch.nn.Linear(num_features, num_features, device=device, bias=False)
+        self.rank = dist.get_global_rank()
+        self.iter = 0
+
+    def forward(self, x):
+        x = self.fc1(x)
+        if self.iter == 3 and x.shape[0] >= 64:
+            raise RuntimeError('CUDA out of memory')
+        x = self.fc2(x)
+        x = self.fc3(x)
+        self.iter += 1
+        return x
+
+    def loss(self, outputs, batch):
+        return torch.sum(outputs)
+
+
+@pytest.mark.gpu
+@pytest.mark.filterwarnings("ignore:`device_train_microbatch_size='auto'` may potentially fail with unexpected.*")
+@pytest.mark.filterwarnings('ignore:CUDA out of memory*')
+@world_size(2)
+def test_fsdp_automicrobatching_sync_hooks(world_size: int):
+    model = SimpleMLPForTestingHooks()
+    model.fc1._fsdp_wrap = True  # pyright: ignore[reportGeneralTypeIssues]
+    model.fc2._fsdp_wrap = True  # pyright: ignore[reportGeneralTypeIssues]
+    dataset = SimpleDatasetForAuto(size=256, feature_size=128)
+    train_dataloader = DataLoader(dataset, batch_size=64, sampler=dist.get_sampler(dataset))
+
+    with patch('composer.trainer.trainer._readd_fsdp_sync_hooks') as mock_readd_hooks:
+        trainer = Trainer(
+            model=model,
+            train_dataloader=train_dataloader,
+            fsdp_config={
+                'forward_prefetch_limit': 1,
+                'backward_prefetch_limit': 1,
+            },
+            max_duration='4ba',
+            device='gpu',
+            device_train_microbatch_size='auto',
+            dist_timeout=20,
+        )
+        trainer.fit()
+
+        # OOM occurs during the 4th batch, so check that sync hooks were readded at the end
+        mock_readd_hooks.assert_called_once()
+
+
 @pytest.mark.gpu
 @world_size(2)
 @pytest.mark.filterwarnings('ignore:Instantiating FSDP with custom process groups.*:UserWarning')
@@ -234,6 +360,62 @@ def test_fsdp_process_group(world_size: int):
     trainer.fit()
 
 
+@pytest.mark.gpu
+@world_size(2)
+@pytest.mark.skipif(
+    version.parse(torch.__version__) < version.parse('2'),
+    reason='FSDP use_orig_params requires torch 2.0 or higher',
+)
+def test_fsdp_subset_of_params_in_opt(world_size: int):
+    model = SimpleModel()
+    dataset = RandomClassificationDataset(size=10)
+    dataloader = DataLoader(dataset, sampler=dist.get_sampler(dataset))
+    optimizer = torch.optim.SGD(model.fc1.parameters(), lr=0.01)
+    unwrapped_optimizer = copy.deepcopy(optimizer)
+
+    trainer = Trainer(
+        model=model,
+        optimizers=optimizer,
+        train_dataloader=dataloader,
+        parallelism_config={
+            'fsdp': {
+                'use_orig_params': True,
+            },
+        },
+        max_duration='3ba',
+    )
+
+    with trainer.state.model.module.summon_full_params(trainer.state.model.module):
+        nb_parameters_before_fsdp = len(unwrapped_optimizer.param_groups[0]['params'])
+        nb_parameters_after_fsdp = len(trainer.state.optimizers[0].param_groups[0]['params'])
+
+        assert nb_parameters_before_fsdp == nb_parameters_after_fsdp
+
+
+@pytest.mark.gpu
+@world_size(2)
+def test_fsdp_subset_of_params_in_opt_without_orig_params(world_size: int):
+    model = SimpleModel()
+    dataset = RandomClassificationDataset(size=10)
+    dataloader = DataLoader(dataset, sampler=dist.get_sampler(dataset))
+    optimizer = torch.optim.SGD(model.fc1.parameters(), lr=0.01)
+
+    expected_error = 'Passing in a subset of model parameters to the optimizer is not supported with use_orig_params=False.'
+
+    with pytest.raises(ValueError, match=expected_error):
+        _ = Trainer(
+            model=model,
+            optimizers=optimizer,
+            train_dataloader=dataloader,
+            parallelism_config={
+                'fsdp': {
+                    'use_orig_params': False,
+                },
+            },
+            max_duration='3ba',
+        )
+
+
 class SimpleMLP(ComposerModel):
 
     def __init__(self, num_features: int = 128, device: str = 'cuda'):
@@ -439,6 +621,28 @@ def test_fsdp_shard(world_size: int):
     )
 
 
+@pytest.mark.gpu
+@world_size(2)
+def test_fsdp_invalid_config_throws_error(world_size: int):
+    model = SimpleModel()
+    model.fc1._fsdp_wrap = True  # pyright: ignore[reportGeneralTypeIssues]
+    model.fc2._fsdp_wrap = True  # pyright: ignore[reportGeneralTypeIssues]
+
+    expected_error = 'activation_cpu_offload=True is not supported with use_orig_params=False.'
+
+    with pytest.raises(ValueError, match=expected_error):
+        _ = Trainer(
+            model=model,
+            parallelism_config={
+                'fsdp': {
+                    'use_orig_params': False,
+                    'activation_cpu_offload': True,
+                },
+            },
+            max_duration='3ba',
+        )
+
+
 @pytest.mark.gpu
 @world_size(2)
 def test_fsdp_shard_and_replicate(world_size: int):
diff --git a/tests/trainer/test_fsdp_param_groups.py b/tests/trainer/test_fsdp_param_groups.py
index 7cbd52520e..8315ec4be2 100644
--- a/tests/trainer/test_fsdp_param_groups.py
+++ b/tests/trainer/test_fsdp_param_groups.py
@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import copy
-import gc
 
 import pytest
 import torch
@@ -11,7 +10,7 @@
 
 from composer.trainer.trainer import Trainer
 from composer.utils import dist, misc
-from tests.common import RandomClassificationDataset, SimpleModel, device, world_size
+from tests.common import EmbeddedWeightTiedModel, RandomClassificationDataset, SimpleModel, device, world_size
 
 
 @pytest.mark.parametrize('mixed_precision', ['DEFAULT'])
@@ -23,7 +22,7 @@ def test_fsdp_param_groups_without_orig_params(mixed_precision: str, device: str
     # Ensure that FSDP with 'use_orig_params=False' raises an exception when passing in an optimizer
     # with multiple param groups
     num_classes = 10
-    model = SimpleModel(num_features=1, num_classes=num_classes)
+    model = SimpleModel(num_features=2, num_classes=num_classes)
     dataset = RandomClassificationDataset(shape=(num_classes,), size=2, num_classes=num_classes)
     dataloader = DataLoader(dataset, sampler=dist.get_sampler(dataset))
 
@@ -31,7 +30,7 @@ def test_fsdp_param_groups_without_orig_params(mixed_precision: str, device: str
     param_groups = [{'params': param, 'lr': (0.1 + 0.1 * i)} for i, param in enumerate(model.parameters())]
     optimizer = torch.optim.SGD(param_groups, lr=0)
 
-    expected_error = 'Multiple optimizer groups with FSDP are only supported with use_orig_params=True.'
+    expected_error = 'Multiple optimizer groups with FSDP are not supported with use_orig_params=False.'
 
     with pytest.raises(RuntimeError, match=expected_error):
         _ = Trainer(
@@ -48,7 +47,6 @@ def test_fsdp_param_groups_without_orig_params(mixed_precision: str, device: str
             max_duration='3ba',
             device=device,
         )
-    gc.collect()
 
 
 @pytest.mark.parametrize('mixed_precision', ['FULL', 'DEFAULT', 'PURE'])
@@ -122,3 +120,86 @@ def test_fsdp_with_param_groups(mixed_precision: str, device: str, reentrant: bo
             assert id(unwrapped_param) != id(wrapped_param)
 
             assert unwrapped_param_group['lr'] == wrapped_param_group['lr']
+
+
+@pytest.mark.parametrize('mixed_precision', ['FULL', 'DEFAULT', 'PURE'])
+@pytest.mark.parametrize('reentrant', [True, False])
+@pytest.mark.filterwarnings('ignore::UserWarning')
+@device('gpu')
+@world_size(2)
+@pytest.mark.skipif(
+    version.parse(torch.__version__) < version.parse('2'),
+    reason='FSDP use_orig_params requires torch 2.0 or higher',
+)
+def test_fsdp_with_param_groups_with_subset_of_params_in_opt(
+    mixed_precision: str,
+    device: str,
+    reentrant: bool,
+    world_size: int,
+):
+    """
+    Test whether an optimizer with param groups and a subset of model variables in the param groups is correctly fsdp wrapped.
+    """
+    num_classes = 10
+
+    # Note that the EmbeddedWeightTiedModel is used instead of SimpleModel to ensure that some of the model parameters
+    # are excluded from the optimzier
+    model = EmbeddedWeightTiedModel(num_features=num_classes)
+    dataset = RandomClassificationDataset(shape=(num_classes,), size=2, num_classes=num_classes)
+    dataloader = DataLoader(dataset, sampler=dist.get_sampler(dataset))
+
+    # create a different group per parameter
+    param_groups = [{
+        'params': model.net1.fc1.parameters(),
+        'lr': 0.1,
+    }, {
+        'params': model.net2.fc2.parameters(),
+        'lr': 0.5,
+    }]
+
+    optimizer = torch.optim.SGD(param_groups)
+    unwrapped_optimizer = copy.deepcopy(optimizer)
+
+    optimizer_groups_pre_fsdp = optimizer.param_groups
+
+    trainer = Trainer(
+        model=model,
+        optimizers=optimizer,
+        train_dataloader=dataloader,
+        parallelism_config={
+            'fsdp': {
+                'activation_checkpointing_reentrant': reentrant,
+                'mixed_precision': mixed_precision,
+            },
+        },
+        max_duration='3ba',
+        device=device,
+    )
+    trainer.fit()
+
+    assert misc.is_model_fsdp(trainer.state.model)
+    trainer_optimizer = trainer.state.optimizers[0]
+    assert len(trainer_optimizer.param_groups) > 1
+    assert len(trainer_optimizer.param_groups) == len(optimizer_groups_pre_fsdp)
+
+    with trainer.state.model.module.summon_full_params(trainer.state.model.module):  # type: ignore
+        for unwrapped_param_group, wrapped_param_group in zip(
+            unwrapped_optimizer.param_groups,
+            trainer_optimizer.param_groups,
+        ):
+
+            unwrapped_param_list = unwrapped_param_group['params']
+            wrapped_param_list = wrapped_param_group['params']
+
+            assert len(unwrapped_param_list) == 1
+            assert len(wrapped_param_list) == 1
+
+            unwrapped_param = unwrapped_param_list[0]
+            wrapped_param = wrapped_param_list[0]
+
+            assert unwrapped_param.shape == wrapped_param.shape
+
+            # the underlying tensor is different because it has been recreated when FSDP wraps the model
+            assert id(unwrapped_param) != id(wrapped_param)
+
+            assert unwrapped_param_group['lr'] == wrapped_param_group['lr']
diff --git a/tests/trainer/test_tp.py b/tests/trainer/test_tp.py
index bfee2e13c9..95b84d32c6 100644
--- a/tests/trainer/test_tp.py
+++ b/tests/trainer/test_tp.py
@@ -47,3 +47,78 @@ def test_tp_train(world_size: int):
     )
 
     trainer.fit()
+
+
+@pytest.mark.gpu
+@world_size(4)
+@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.3'), reason='requires PyTorch 2.3+')
+@pytest.mark.filterwarnings(r'ignore:.*\(TP\) is experimental.*:FutureWarning')
+def test_tp_with_param_groups(world_size: int):
+    from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel
+
+    model = SimpleModel()
+    dataset = RandomClassificationDataset(size=8)
+    dataloader = DataLoader(dataset, batch_size=2, sampler=dist.get_sampler(dataset))
+    optimizer = torch.optim.SGD([{
+        'params': model.fc1.parameters(),
+        'lr': 0.1,
+    }, {
+        'params': model.fc2.parameters(),
+        'lr': 0.5,
+    }])
+
+    layer_plan = {
+        'fc1': ColwiseParallel(),
+        'fc2': RowwiseParallel(),
+    }
+
+    expected_error = 'Multiple optimizer groups are not supported with tensor parallelism.'
+
+    with pytest.raises(RuntimeError, match=expected_error):
+        _ = Trainer(
+            model=model,
+            optimizers=optimizer,
+            train_dataloader=dataloader,
+            parallelism_config={
+                'tp': {
+                    'layer_plan': layer_plan,
+                    'tensor_parallel_degree': 2,
+                },
+                'fsdp': {},
+            },
+            max_duration='3ba',
+        )
+
+
+@pytest.mark.gpu
+@world_size(4)
+@pytest.mark.skipif(version.parse(torch.__version__) < version.parse('2.3'), reason='requires PyTorch 2.3+')
+@pytest.mark.filterwarnings(r'ignore:.*\(TP\) is experimental.*:FutureWarning')
+def test_tp_with_subset_of_params(world_size: int):
+    from torch.distributed.tensor.parallel import ColwiseParallel
+
+    model = SimpleModel()
+    dataset = RandomClassificationDataset(size=8)
+    dataloader = DataLoader(dataset, batch_size=2, sampler=dist.get_sampler(dataset))
+    optimizer = torch.optim.SGD(model.fc1.parameters(), lr=0.1)
+
+    layer_plan = {
+        'fc1': ColwiseParallel(),
+    }
+
+    expected_error = 'Passing in a subset of model parameters to the optimizer is not supported with tensor parallelism.'
+
+    with pytest.raises(ValueError, match=expected_error):
+        _ = Trainer(
+            model=model,
+            optimizers=optimizer,
+            train_dataloader=dataloader,
+            parallelism_config={
+                'tp': {
+                    'layer_plan': layer_plan,
+                    'tensor_parallel_degree': 2,
+                },
+                'fsdp': {},
+            },
+            max_duration='3ba',
+        )
diff --git a/tests/utils/object_store/test_azure_object_store.py b/tests/utils/object_store/test_azure_object_store.py
index 949e2149ff..3a83ee6a74 100644
--- a/tests/utils/object_store/test_azure_object_store.py
+++ b/tests/utils/object_store/test_azure_object_store.py
@@ -9,6 +9,7 @@
 
 
 @pytest.mark.remote
+@pytest.mark.skip(reason='Waiting for new Azure key to be approved')
 def test_azure_object_store_integration():
     model = SimpleModel()
     train_dataloader = DataLoader(dataset=RandomClassificationDataset())
diff --git a/tests/utils/object_store/test_gs_object_store.py b/tests/utils/object_store/test_gs_object_store.py
index 006fd933bc..93b8a05ccd 100644
--- a/tests/utils/object_store/test_gs_object_store.py
+++ b/tests/utils/object_store/test_gs_object_store.py
@@ -9,7 +9,6 @@
 from botocore.exceptions import ClientError
 from torch.utils.data import DataLoader
 
-from composer.loggers import RemoteUploaderDownloader
 from composer.optim import DecoupledSGDW
 from composer.trainer import Trainer
 from composer.utils import GCSObjectStore
@@ -17,8 +16,9 @@
 
 
 def get_gcs_os_from_trainer(trainer: Trainer) -> GCSObjectStore:
-    rud = [dest for dest in trainer.logger.destinations if isinstance(dest, RemoteUploaderDownloader)][0]
-    gcs_os = rud.remote_backend
+    assert trainer._checkpoint_saver is not None
+    assert trainer._checkpoint_saver.remote_uploader is not None
+    gcs_os = trainer._checkpoint_saver.remote_uploader.remote_backend
     assert isinstance(gcs_os, GCSObjectStore)
     return gcs_os
 
@@ -34,7 +34,7 @@ def test_gs_object_store_integration_hmac_auth(expected_use_gcs_sdk_val=False, c
         model=model,
         optimizers=optimizer,
         train_dataloader=train_dataloader,
-        save_folder='gs://mosaicml-internal-integration-testing/checkpoints/{run_name}',
+        save_folder='gs://mosaicml-runtime-internal-integration-testing/checkpoints/{run_name}',
         save_filename='test-model.pt',
         max_duration='1ba',
         precision='amp_bf16',
@@ -53,7 +53,7 @@ def test_gs_object_store_integration_hmac_auth(expected_use_gcs_sdk_val=False, c
         model=model,
         optimizers=optimizer,
         train_dataloader=train_dataloader,
-        load_path=f'gs://mosaicml-internal-integration-testing/checkpoints/{run_name}/test-model.pt',
+        load_path=f'gs://mosaicml-runtime-internal-integration-testing/checkpoints/{run_name}/test-model.pt',
         max_duration='2ba',
         precision='amp_bf16',
     )