diff --git a/llmfoundry/callbacks/curriculum_learning_callback.py b/llmfoundry/callbacks/curriculum_learning_callback.py
index 961bf1cae1..43a194b04e 100644
--- a/llmfoundry/callbacks/curriculum_learning_callback.py
+++ b/llmfoundry/callbacks/curriculum_learning_callback.py
@@ -8,22 +8,28 @@
 """
 
 import logging
-from typing import Any, Dict
+import copy
+from typing import Any
 
-from composer.core import State
-from composer.loggers import Logger
+from composer import DataSpec
+from composer.core import State, TimeUnit, ensure_time
+from composer.loggers import Logger, MosaicMLLogger
+from composer.trainer.trainer import _get_initial_device_train_microbatch_size
 from streaming import StreamingDataset
+from streaming.base.util import clean_stale_shared_memory
 from torch.utils.data import DataLoader
 
+from llmfoundry.data.dataloader import build_dataloader
 from llmfoundry.interfaces import CallbackWithConfig
-from llmfoundry.utils.warnings import experimental_class
+from llmfoundry.utils.builders import build_tokenizer
+from llmfoundry.utils.config_utils import calculate_batch_size_info
+from llmfoundry.utils.exceptions import BaseContextualError, TrainDataLoaderLocation
 
 log = logging.getLogger(__name__)
 
 __all__ = ['CurriculumLearning']
 
 
-@experimental_class('CurriculumLearning callback')
 class CurriculumLearning(CallbackWithConfig):
     """Starts an epoch with a different dataset when resuming from a checkpoint.
 
@@ -34,20 +40,179 @@ class CurriculumLearning(CallbackWithConfig):
         dataset_index (int): The index of the dataset currently being used.
     """
 
-    def __init__(self, train_config: Dict, dataset_index: int):
-        self.dataset_index = dataset_index
-        self.saved_dataset_index = 0
-        self.all_dataset_configs = []
-        self.current_dataset_state = {}
-        # The current dataset config is resolved and passed in train.py
-        self.current_dataset_config = train_config['train_loader']
+    def __init__(
+        self, train_config: dict[str, Any], duration: str | int | TimeUnit,
+        schedule: list[dict[str, Any]]
+    ):
+        non_positive_error = ValueError('The duration must be positive.')
+        unit_error = ValueError(
+            'Schedules can only be defined in terms of epochs or tokens.'
+        )
+
+        # Ensure all duration values are positive
+        # Ensure all duration units are in epochs or  tokens
+        self._duration = ensure_time(duration, TimeUnit.EPOCH)
+        if self._duration.value <= 0:
+            raise non_positive_error
+        if self._duration.unit != TimeUnit.EPOCH and self._duration.unit != TimeUnit.TOKEN:
+            raise unit_error
+
+        self._schedule = schedule
+        for datamix in self._schedule:
+            assert 'duration' in datamix, 'Each datamix must have a duration.'
+            datamix['duration'] = ensure_time(
+                datamix['duration'], TimeUnit.EPOCH
+            )
+            if datamix['duration'].value <= 0:
+                raise non_positive_error
+            if datamix['duration'].unit != TimeUnit.EPOCH and datamix[
+                'duration'].unit != TimeUnit.TOKEN:
+                raise unit_error
+            assert 'train_loader' in datamix, 'Each datamix must have a train_loader.'
+
+        self._schedule_index = -1
+
+        # Copied from llmfoundry/utils/config_utils.py
+        self.device_train_batch_size, _, _ = calculate_batch_size_info(
+            train_config['global_train_batch_size'],
+            train_config['device_train_microbatch_size'],
+            data_replication_degree=1,
+        )
+
+        # Copied from scripts/train/train.py
+        tokenizer_name = train_config['tokenizer']['name']
+        tokenizer_kwargs = train_config['tokenizer'].get('kwargs', {})
+        self.tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
 
     def before_load(self, state: State, logger: Logger):
         del logger
 
-        # Save the current dataset state so we can restore it correctly
-        # if we are resuming with a new dataset.
-        train_loader = state.train_dataloader
+        # Ensure all duration units are the same as max_duration
+        units_match = True
+        if self._duration.unit != state.max_duration.unit:
+            units_match = False
+        for datamix in self._schedule:
+            if datamix['duration'].unit != state.max_duration.unit:
+                units_match = False
+        if not units_match:
+            raise ValueError((
+                'All durations in the schedule must have the same units as '
+                'the max_duration.'
+            ))
+
+        # Ensure schedule duration is greater than max_duration
+        schedule_duration = self._duration
+        for datamix in self._schedule:
+            schedule_duration += datamix['duration']
+        if schedule_duration < state.max_duration:
+            raise ValueError((
+                'The sum of all durations in the schedule must be greater than '
+                'or equal to the max_duration.'
+            ))
+
+        self._validate_dataloader(state.train_dataloader)
+
+    def after_load(self, state: State, logger: Logger):
+        del logger
+
+        self._validate_dataloader(state.train_dataloader)
+
+        # Check if adding a new datamix to a run that didn't use this callback
+        if self._schedule_index == -1 and state.timestamp >= self._duration:
+            self._schedule_index = 0
+            state.timestamp = state.timestamp.to_next_iteration()
+        # If checkpoint was saved before iteration was incremented, we need to increment it now
+        elif ((
+            self._schedule[self._schedule_index]['duration'].unit
+            == TimeUnit.TOKEN and state.timestamp.token_in_iteration
+            >= self._schedule[self._schedule_index]['duration'].value
+        ) or (
+            self._schedule[self._schedule_index]['duration'].unit
+            == TimeUnit.EPOCH and state.timestamp.epoch_in_iteration
+            >= self._schedule[self._schedule_index]['duration'].value
+        )):
+            log.warning((
+                'The CurriculumLearning callback has detected that the previous run did not correctly '
+                'increment the iteration.'
+            ))
+            self._schedule_index += 1
+            state.timestamp = state.timestamp.to_next_iteration()
+
+    def iteration_start(self, state: State, logger: Logger):
+        # Reset and initialize state train dataloader
+        log.warning(
+            'trainer._train_data_spec should be updated whenever the dataloader is updated'
+        )
+
+        # Swap the dataset if starting a new iteration that's not the original datamix
+        if self._schedule_index >= 0:
+            clean_stale_shared_memory()
+            data_spec = self._build_train_loader(
+                train_loader_config=copy.deepcopy(
+                    self._schedule[self._schedule_index]
+                )['train_loader'],
+                logger=logger,
+            )
+            state.set_dataloader(
+                dataloader=data_spec.dataloader,
+                dataloader_label='train',
+            )
+            # state.train_dataloader = state.dataloader
+            state.device_train_microbatch_size = _get_initial_device_train_microbatch_size(
+                state.device_train_microbatch_size,
+                state.auto_microbatching,
+                state.train_dataloader,
+            )
+            self._validate_dataloader(state.train_dataloader)
+
+        # Set the length of the new iteration
+        if self._schedule_index == -1:
+            state._iteration_length = self._duration
+        else:
+            state._iteration_length = self._schedule[self._schedule_index
+                                                    ]['duration']
+
+    def iteration_end(self, state: State, logger: Logger):
+        del state, logger  # unused
+
+        self._schedule_index += 1
+
+    def state_dict(self):
+        return {
+            'duration': self._duration,
+            'schedule': self._schedule,
+            'schedule_index': self._schedule_index,
+        }
+
+    def load_state_dict(self, state: dict[str, Any]):
+        # Ensure that the schedule has not changed on already trained datamixes
+        assert self._duration == state['duration']
+        for idx in range(state['schedule_index'] + 1):
+            assert self._schedule[idx] == state['schedule'][idx]
+
+        self._schedule_index = state['schedule_index']
+
+    def _build_train_loader(
+        self, train_loader_config: dict[str, Any], logger: Logger
+    ) -> DataSpec:
+        # Copied from scripts/train/train.py
+        log.info(
+            f'Building train loader in CurriculumLearning callback for dataset {self._schedule_index}'
+        )
+        try:
+            return build_dataloader(
+                train_loader_config,
+                self.tokenizer,
+                self.device_train_batch_size,
+            )
+        except BaseContextualError as e:
+            for destination in logger.destinations:
+                if isinstance(destination, MosaicMLLogger):
+                    e.location = TrainDataLoaderLocation
+                    destination.log_exception(e)
+            raise e
+
+    def _validate_dataloader(self, train_loader: Any):
         # Check if we are using a DataLoader and StreamingDataset
         if not isinstance(train_loader, DataLoader):
             raise ValueError(
@@ -61,54 +226,3 @@ def before_load(self, state: State, logger: Logger):
                 f'because it requires loading and saving dataset state. ',
                 f'Instead, got a dataset of type {type(dataset)}',
             )
-        assert isinstance(dataset, StreamingDataset)
-        # Save the current dataset state so we can restore it if needed.
-        self.current_dataset_state = dataset.state_dict(  # type: ignore
-            num_samples=0, from_beginning=False)
-
-    def after_load(self, state: State, logger: Logger):
-        del logger
-
-        # As saved_dataset_index is loaded from state_dict, this only runs when
-        # a user explicitly increments the dataset_index and not on any other
-        # resumption, including autoresume.
-        train_loader = state._train_dataloader
-        assert isinstance(
-            train_loader,
-            DataLoader,
-        ), 'CurriculumLearning callback requires a DataLoader.'
-        dataset = train_loader.dataset
-        assert isinstance(
-            dataset,
-            StreamingDataset,
-        ), 'CurriculumLearning callback requires a StreamingDataset.'
-        if self.saved_dataset_index < self.dataset_index:
-            # Ignore the dataset state that was read in from the checkpoint, and
-            # replace with the new dataset state. This preserves resumption info.
-            if self.current_dataset_state['epoch'] < 0:
-                # Make sure the epoch in the loaded state dict is not negative.
-                # Since `__iter__` has not yet been called on the dataset, the
-                # epoch index in the dataset will still be -1. We need to ensure
-                # that we set the epoch correctly to 0 in this case.
-                self.current_dataset_state['epoch'] = 0
-            dataset.load_state_dict(  # type: ignore
-                self.current_dataset_state)
-            # Start a new epoch since we are using a new dataset.
-            # This will also reset the sample_in_epoch written to checkpoint,
-            # making sure that subsequent resumptions proceed correctly.
-            state.timestamp = state.timestamp.to_next_epoch()
-            # Append the new dataset config to the list of all dataset configs.
-            self.all_dataset_configs.append(self.current_dataset_config)
-        elif self.dataset_index == 0 and len(self.all_dataset_configs) == 0:
-            # Make sure to track our current dataset config if we are just starting training.
-            self.all_dataset_configs.append(self.current_dataset_config)
-
-    def state_dict(self):
-        return {
-            'dataset_index': self.dataset_index,
-            'all_dataset_configs': self.all_dataset_configs,
-        }
-
-    def load_state_dict(self, state: Dict[str, Any]):
-        self.saved_dataset_index = state.get('dataset_index', 0)
-        self.all_dataset_configs = state.get('all_dataset_configs', [])
diff --git a/llmfoundry/data/text_data.py b/llmfoundry/data/text_data.py
index 86d5edbaf4..db64a397e5 100644
--- a/llmfoundry/data/text_data.py
+++ b/llmfoundry/data/text_data.py
@@ -289,7 +289,7 @@ def get_sequence_id_from_batch(
         return torch.cat([left_zeros, cumulative_sep[:, :-1]], dim=1)
 
 
-def build_streams(streams: Optional[Dict[str, Any]] = None,):
+def build_streams(streams: Optional[Dict[str, Any]] = None):
     streams_dict = streams
     # build streams
     streams_ret = []
diff --git a/llmfoundry/utils/__init__.py b/llmfoundry/utils/__init__.py
index dd43efcdd7..1752e024e7 100644
--- a/llmfoundry/utils/__init__.py
+++ b/llmfoundry/utils/__init__.py
@@ -2,9 +2,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from llmfoundry.utils.builders import (
+    add_metrics_to_eval_loaders,
     build_algorithm,
-    build_callback,
     build_composer_model,
+    build_eval_loaders,
     build_evaluators,
     build_icl_data_and_gauntlet,
     build_icl_evaluators,
@@ -12,18 +13,18 @@
     build_metric,
     build_optimizer,
     build_scheduler,
-    build_tokenizer,
 )
+from llmfoundry.utils.callback_builder import build_callback
 from llmfoundry.utils.checkpoint_conversion_helpers import (
     convert_and_save_ft_weights,
     get_hf_tokenizer_from_composer_state_dict,
     load_tokenizer,
 )
+from llmfoundry.utils.config_dist_utils import process_init_device
 from llmfoundry.utils.config_utils import (
     calculate_batch_size_info,
     log_config,
     pop_config,
-    process_init_device,
     update_batch_size_info,
 )
 from llmfoundry.utils.data_prep_utils import (
@@ -60,8 +61,10 @@
 )
 
 __all__ = [
+    'add_metrics_to_eval_loaders',
     'build_algorithm',
     'build_callback',
+    'build_eval_loaders',
     'build_evaluators',
     'build_icl_data_and_gauntlet',
     'build_icl_evaluators',
diff --git a/llmfoundry/utils/builders.py b/llmfoundry/utils/builders.py
index 73eb026d98..9f27a9923e 100644
--- a/llmfoundry/utils/builders.py
+++ b/llmfoundry/utils/builders.py
@@ -19,7 +19,7 @@
 )
 
 import torch
-from composer.core import Algorithm, Callback, Evaluator
+from composer.core import Algorithm, Evaluator
 from composer.loggers import LoggerDestination
 from composer.models import ComposerModel
 from composer.optim.scheduler import ComposerScheduler
@@ -31,7 +31,7 @@
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
 from llmfoundry import registry
-from llmfoundry.callbacks import EvalGauntlet
+from llmfoundry.callbacks.eval_gauntlet_callback import EvalGauntlet
 from llmfoundry.data.dataloader import build_dataloader
 from llmfoundry.eval.datasets.in_context_learning_evaluation import \
     get_icl_task_dataloader
@@ -43,7 +43,6 @@
 
 __all__ = [
     'build_algorithm',
-    'build_callback',
     'build_evaluators',
     'build_icl_data_and_gauntlet',
     'build_icl_evaluators',
@@ -234,33 +233,6 @@ def build_composer_model(
     return model
 
 
-def build_callback(
-    name: str,
-    kwargs: Optional[Dict[str, Any]] = None,
-    train_config: Any = None,
-) -> Callback:
-    """Builds a callback from the registry."""
-    registry_to_use = registry.callbacks
-    if name in registry.callbacks_with_config:
-        if kwargs is None:
-            kwargs = {}
-        if 'train_config' in kwargs:
-            raise ValueError(
-                f'`train_config` is a reserved keyword for callbacks with config. Please remove it from the kwargs.',
-            )
-        kwargs['train_config'] = train_config
-        registry_to_use = registry.callbacks_with_config
-
-    return construct_from_registry(
-        name=name,
-        registry=registry_to_use,
-        partial_function=True,
-        pre_validation_function=Callback,
-        post_validation_function=None,
-        kwargs=kwargs,
-    )
-
-
 def build_logger(
     name: str,
     kwargs: Optional[Dict[str, Any]] = None,
@@ -455,8 +427,8 @@ def build_tokenizer(
 
     signal_file_path = f'.node_{dist.get_node_rank()}_local_rank0_completed_tokenizer_setup'
 
-    if dist.is_available() and dist.is_initialized(
-    ) and dist.get_world_size() > 1:
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size(
+    ) > 1:
         # Make sure the tokenizer files are downloaded and cached first by local rank 0
         with dist.local_rank_zero_download_and_wait(signal_file_path):
             pass
@@ -482,8 +454,8 @@ def build_tokenizer(
             f'The tokenizer {tokenizer_name} must have an eos_token.',
         )
 
-    if dist.is_available() and dist.is_initialized(
-    ) and dist.get_world_size() > 1:
+    if dist.is_available() and dist.is_initialized() and dist.get_world_size(
+    ) > 1:
         if dist.get_local_rank() == 0:
             with open(signal_file_path, 'wb') as f:
                 f.write(b'local_rank0_completed_tokenizer_setup')
diff --git a/llmfoundry/utils/callback_builder.py b/llmfoundry/utils/callback_builder.py
new file mode 100644
index 0000000000..b0387c6a9f
--- /dev/null
+++ b/llmfoundry/utils/callback_builder.py
@@ -0,0 +1,36 @@
+# Copyright 2022 MosaicML LLM Foundry authors
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Optional
+
+from composer.core import Callback
+
+from llmfoundry import registry
+from llmfoundry.utils.registry_utils import construct_from_registry
+
+
+def build_callback(
+    name: str,
+    kwargs: Optional[dict[str, Any]] = None,
+    train_config: Any = None,
+) -> Callback:
+    """Builds a callback from the registry."""
+    registry_to_use = registry.callbacks
+    if name in registry.callbacks_with_config:
+        if kwargs is None:
+            kwargs = {}
+        if 'train_config' in kwargs:
+            raise ValueError(
+                f'`train_config` is a reserved keyword for callbacks with config. Please remove it from the kwargs.',
+            )
+        kwargs['train_config'] = train_config
+        registry_to_use = registry.callbacks_with_config
+
+    return construct_from_registry(
+        name=name,
+        registry=registry_to_use,
+        partial_function=True,
+        pre_validation_function=Callback,
+        post_validation_function=None,
+        kwargs=kwargs,
+    )
diff --git a/llmfoundry/utils/config_dist_utils.py b/llmfoundry/utils/config_dist_utils.py
new file mode 100644
index 0000000000..f75821aae8
--- /dev/null
+++ b/llmfoundry/utils/config_dist_utils.py
@@ -0,0 +1,81 @@
+import contextlib
+import warnings
+from typing import Any, Mapping, Optional
+
+from llmfoundry.layers_registry import ffns_with_megablocks
+from llmfoundry.models.utils import init_empty_weights
+
+__all__ = ['process_init_device']
+
+
+def process_init_device(model_cfg: dict[str, Any], fsdp_config: Optional[dict]):
+    # Restrict model init_device to 'meta' and 'cpu',
+    # using 'cuda' vs. 'cuda:id' is tricky and can lead to common user errors
+    # when multiple GPUs are available.
+    # Also 'meta' is only valid when using FSDP
+    init_context = contextlib.nullcontext()
+    if 'init_device' in model_cfg:
+        assert model_cfg['init_device'] in ['meta', 'cpu', 'mixed']
+        if fsdp_config is None and model_cfg['init_device'] == 'meta':
+            warnings.warn(
+                "Using `cfg.model.init_device='meta'` is only valid when using FSDP! " +\
+                "Reverting to `cfg.model.init_device='cpu'`.")
+            model_cfg['init_device'] = 'cpu'
+        if model_cfg['init_device'] == 'meta':
+            init_context = init_empty_weights()
+        if model_cfg['init_device'] == 'mixed':
+            if fsdp_config is None:
+                raise NotImplementedError(
+                    'Using init_device `mixed` is only supported with FSDP. ' +
+                    'Please add a FSDP config.',
+                )
+            # Always set `sync_module_states` to True for mixed initialization
+            if not fsdp_config.get('sync_module_states', False):
+                warnings.warn((
+                    'Setting `sync_module_states = True` for FSDP. This is required '
+                    'when using mixed initialization.'
+                ))
+                fsdp_config['sync_module_states'] = True
+
+            # Set defaults for mixed initialization
+            fsdp_config.setdefault('use_orig_params', False)
+            fsdp_config.setdefault('load_monolith_rank0_only', True)
+
+    # Set ffn_config.device_mesh to fsdp_config.device_mesh
+    if fsdp_config is not None and 'device_mesh' in fsdp_config and 'ffn_config' in model_cfg and model_cfg[
+        'ffn_config'].get('ffn_type', None) in ffns_with_megablocks:
+        # Raise ValueError if not using device mesh with MoE expert parallelism
+        if fsdp_config['device_mesh'] is None and model_cfg['ffn_config'].get(
+            'moe_world_size',
+            1,
+        ) > 1:
+            raise ValueError(
+                'device_mesh must be specified in fsdp_config when using MoE with moe_world_size > 1.',
+            )
+        model_cfg['ffn_config']['device_mesh'] = fsdp_config['device_mesh']
+
+    # No mixed precision needed for weights when they're already 16 bits
+    master_dtype = model_cfg.get('master_weights_dtype')
+    small_dtypes = (
+        'bf16',
+        'fp16',
+        'float16',
+        'bfloat16',
+        'amp_fp16',
+        'amp_bf16',
+    )
+    if fsdp_config and master_dtype in small_dtypes:
+        reduce_dtype = None
+        buffer_dtype = None
+        mixed_precision = fsdp_config.get('mixed_precision')
+        if isinstance(mixed_precision, Mapping):
+            reduce_dtype = mixed_precision.get('reduce_dtype')
+            buffer_dtype = mixed_precision.get('buffer_dtype')
+        fsdp_config['mixed_precision'] = {
+            'param_dtype': None,
+            'reduce_dtype': reduce_dtype,
+            'buffer_dtype': buffer_dtype,
+            'keep_low_precision_grads': True,
+        }
+
+    return init_context
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
index 5c1ec9114a..250a59bba9 100644
--- a/llmfoundry/utils/config_utils.py
+++ b/llmfoundry/utils/config_utils.py
@@ -1,7 +1,6 @@
 # Copyright 2022 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-import contextlib
 import copy
 import logging
 import math
@@ -14,7 +13,6 @@
     Dict,
     List,
     Literal,
-    Mapping,
     Optional,
     Set,
     Tuple,
@@ -28,16 +26,12 @@
 from omegaconf import OmegaConf as om
 from transformers import PretrainedConfig
 
-from llmfoundry.layers_registry import ffns_with_megablocks
-from llmfoundry.models.utils import init_empty_weights
-
 log = logging.getLogger(__name__)
 
 __all__ = [
     'pop_config',
     'calculate_batch_size_info',
     'update_batch_size_info',
-    'process_init_device',
     'log_config',
     'log_dataset_uri',
 ]
@@ -423,79 +417,6 @@ def update_batch_size_info(cfg: Dict[str, Any]) -> Dict[str, Any]:
     return cfg
 
 
-def process_init_device(model_cfg: Dict[str, Any], fsdp_config: Optional[Dict]):
-    # Restrict model init_device to 'meta' and 'cpu',
-    # using 'cuda' vs. 'cuda:id' is tricky and can lead to common user errors
-    # when multiple GPUs are available.
-    # Also 'meta' is only valid when using FSDP
-    init_context = contextlib.nullcontext()
-    if 'init_device' in model_cfg:
-        assert model_cfg['init_device'] in ['meta', 'cpu', 'mixed']
-        if fsdp_config is None and model_cfg['init_device'] == 'meta':
-            warnings.warn(
-                "Using `cfg.model.init_device='meta'` is only valid when using FSDP! " +\
-                "Reverting to `cfg.model.init_device='cpu'`.")
-            model_cfg['init_device'] = 'cpu'
-        if model_cfg['init_device'] == 'meta':
-            init_context = init_empty_weights()
-        if model_cfg['init_device'] == 'mixed':
-            if fsdp_config is None:
-                raise NotImplementedError(
-                    'Using init_device `mixed` is only supported with FSDP. ' +
-                    'Please add a FSDP config.',
-                )
-            # Always set `sync_module_states` to True for mixed initialization
-            if not fsdp_config.get('sync_module_states', False):
-                warnings.warn((
-                    'Setting `sync_module_states = True` for FSDP. This is required '
-                    'when using mixed initialization.'
-                ))
-                fsdp_config['sync_module_states'] = True
-
-            # Set defaults for mixed initialization
-            fsdp_config.setdefault('use_orig_params', False)
-            fsdp_config.setdefault('load_monolith_rank0_only', True)
-
-    # Set ffn_config.device_mesh to fsdp_config.device_mesh
-    if fsdp_config is not None and 'device_mesh' in fsdp_config and 'ffn_config' in model_cfg and model_cfg[
-        'ffn_config'].get('ffn_type', None) in ffns_with_megablocks:
-        # Raise ValueError if not using device mesh with MoE expert parallelism
-        if fsdp_config['device_mesh'] is None and model_cfg['ffn_config'].get(
-            'moe_world_size',
-            1,
-        ) > 1:
-            raise ValueError(
-                'device_mesh must be specified in fsdp_config when using MoE with moe_world_size > 1.',
-            )
-        model_cfg['ffn_config']['device_mesh'] = fsdp_config['device_mesh']
-
-    # No mixed precision needed for weights when they're already 16 bits
-    master_dtype = model_cfg.get('master_weights_dtype')
-    small_dtypes = (
-        'bf16',
-        'fp16',
-        'float16',
-        'bfloat16',
-        'amp_fp16',
-        'amp_bf16',
-    )
-    if fsdp_config and master_dtype in small_dtypes:
-        reduce_dtype = None
-        buffer_dtype = None
-        mixed_precision = fsdp_config.get('mixed_precision')
-        if isinstance(mixed_precision, Mapping):
-            reduce_dtype = mixed_precision.get('reduce_dtype')
-            buffer_dtype = mixed_precision.get('buffer_dtype')
-        fsdp_config['mixed_precision'] = {
-            'param_dtype': None,
-            'reduce_dtype': reduce_dtype,
-            'buffer_dtype': buffer_dtype,
-            'keep_low_precision_grads': True,
-        }
-
-    return init_context
-
-
 def log_config(cfg: Dict[str, Any]) -> None:
     """Logs the current config and updates the wandb and mlflow configs.
 
diff --git a/scripts/eval/eval.py b/scripts/eval/eval.py
index 8a0a5c104f..7deb5c4397 100644
--- a/scripts/eval/eval.py
+++ b/scripts/eval/eval.py
@@ -26,18 +26,18 @@
 install()
 from llmfoundry.utils.builders import (
     add_metrics_to_eval_loaders,
-    build_callback,
     build_composer_model,
     build_evaluators,
     build_logger,
     build_tokenizer,
 )
+from llmfoundry.utils.callback_builder import build_callback
+from llmfoundry.utils.config_dist_utils import process_init_device
 from llmfoundry.utils.config_utils import (
     EVAL_CONFIG_KEYS,
     EvalConfig,
     log_config,
     make_dataclass_and_log_config,
-    process_init_device,
 )
 from llmfoundry.utils.registry_utils import import_file
 
diff --git a/scripts/train/train.py b/scripts/train/train.py
index f2a70b526d..bf55d69a8b 100644
--- a/scripts/train/train.py
+++ b/scripts/train/train.py
@@ -34,7 +34,6 @@
 from llmfoundry.utils.builders import (
     add_metrics_to_eval_loaders,
     build_algorithm,
-    build_callback,
     build_composer_model,
     build_evaluators,
     build_logger,
@@ -42,6 +41,8 @@
     build_scheduler,
     build_tokenizer,
 )
+from llmfoundry.utils.callback_builder import build_callback
+from llmfoundry.utils.config_dist_utils import process_init_device
 from llmfoundry.utils.config_utils import (
     TRAIN_CONFIG_KEYS,
     TrainConfig,
@@ -49,7 +50,6 @@
     log_dataset_uri,
     make_dataclass_and_log_config,
     pop_config,
-    process_init_device,
     update_batch_size_info,
 )
 from llmfoundry.utils.exceptions import (
diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
index 0577e13a1f..e4bd1dbe97 100644
--- a/tests/a_scripts/inference/test_convert_composer_to_hf.py
+++ b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -33,7 +33,8 @@
     build_optimizer,
     build_tokenizer,
 )
-from llmfoundry.utils.config_utils import process_init_device, to_dict_container
+from llmfoundry.utils.config_dist_utils import process_init_device
+from llmfoundry.utils.config_utils import to_dict_container
 from scripts.inference.convert_composer_to_hf import convert_composer_to_hf
 from tests.data_utils import make_tiny_ft_dataset
 
diff --git a/tests/callbacks/test_async_eval_callback.py b/tests/callbacks/test_async_eval_callback.py
index e0794df08f..d3974003a3 100644
--- a/tests/callbacks/test_async_eval_callback.py
+++ b/tests/callbacks/test_async_eval_callback.py
@@ -15,7 +15,7 @@
     validate_eval_run_config,
     validate_interval,
 )
-from llmfoundry.utils.builders import build_callback
+from llmfoundry.utils.callback_builder import build_callback
 from mcli import Run, RunConfig, RunStatus
 
 RUN_NAME = 'foo_bar-1234'
diff --git a/tests/callbacks/test_curriculum_learning_callback.py b/tests/callbacks/test_curriculum_learning_callback.py
index bbdbf3d691..5549b90947 100644
--- a/tests/callbacks/test_curriculum_learning_callback.py
+++ b/tests/callbacks/test_curriculum_learning_callback.py
@@ -1,7 +1,7 @@
 # Copyright 2024 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-from llmfoundry.utils.builders import build_callback
+from llmfoundry.utils.callback_builder import build_callback
 
 
 def test_curriculum_learning_callback_builds():
diff --git a/tests/callbacks/test_mbmoe_tok_per_expert_callback.py b/tests/callbacks/test_mbmoe_tok_per_expert_callback.py
index aff9f1724b..29a960e539 100644
--- a/tests/callbacks/test_mbmoe_tok_per_expert_callback.py
+++ b/tests/callbacks/test_mbmoe_tok_per_expert_callback.py
@@ -1,7 +1,7 @@
 # Copyright 2024 MosaicML LLM Foundry authors
 # SPDX-License-Identifier: Apache-2.0
 
-from llmfoundry.utils.builders import build_callback
+from llmfoundry.utils.callback_builder import build_callback
 
 
 def test_mbmoe_tok_per_expert_builds():
diff --git a/tests/utils/test_builders.py b/tests/utils/test_builders.py
index dfcb5b327c..22fa61e6d2 100644
--- a/tests/utils/test_builders.py
+++ b/tests/utils/test_builders.py
@@ -19,13 +19,13 @@
 from llmfoundry.tokenizers.tiktoken import TiktokenTokenizerWrapper
 from llmfoundry.utils.builders import (
     add_metrics_to_eval_loaders,
-    build_callback,
     build_eval_loaders,
     build_evaluators,
     build_logger,
     build_optimizer,
     build_tokenizer,
 )
+from llmfoundry.utils.callback_builder import build_callback
 
 
 @pytest.mark.parametrize(