Supporting automicrobatching on FSDP2

composer/distributed/dist_strategy.py

@@ -31,6 +31,7 @@
     get_mixed_precision,
     set_custom_fsdp_module_kwargs,
 )
+from composer.distributed.shared_utils import add_fsdp_oom_hooks
 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']
@@ -262,22 +263,6 @@ def prepare_fsdp_module(
     # 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
-    # FSPD block with a hook that checks if any other rank OOMed.
-    def sync_hook(*args):
-        # Check if any other rank hit an OOM
-        found_cuda_oom_tensor = 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 = 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('CUDA out of memory encountered on a different rank')
-
     # Necessary variables for optimizers with multiple param groups in FSDP
     param_name_to_group_num = None
     group_num_to_opt_group_info = None
@@ -492,20 +477,8 @@ def lambda_fn(module: torch.nn.Module) -> Union[bool, dict]:
                 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))
+                hook_handles = add_fsdp_oom_hooks(fsdp_obj, fsdp_config_version=1, device=device)
                 fsdp_obj_named_modules.update(dict(fsdp_obj.named_modules()))
 
             if hasattr(fsdp_obj, '_exec_order_data'):

composer/distributed/fsdp2.py

@@ -8,7 +8,8 @@
 
 import torch
 import torch.nn as nn
-from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard
+from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
 from torch.distributed.fsdp.wrap import CustomPolicy
 
 from composer.distributed.fsdp2_utils import (
@@ -17,7 +18,9 @@
     get_standalone_and_tied_modules,
     legalize_param_sharing_between_modules,
 )
-from composer.utils.parallelism import FSDP2Config
+from composer.distributed.shared_utils import add_fsdp_oom_hooks, get_direct_children_from_composer_model
+from composer.models import ComposerModel
+from composer.utils import FSDP2Config
 
 log = logging.getLogger(__name__)
 
@@ -115,16 +118,13 @@ def prepare_fully_shard(
     model: nn.Module,
     fsdp2_config: FSDP2Config,
     auto_wrap_policy: Optional[CustomPolicy] = None,
-) -> None:
+):
     """Applies FSDP2's `fully_shard` to the model according to given fsdp2_config.
 
     Args:
         model (torch.nn.Module): The model to prepare.
         fsdp2_config (FSDP2Config): The FSDP2 configuration.
         auto_wrap_policy (Optional[CustomPolicy]): The policy to apply to the model.
-
-    Returns:
-        None
     """
     # If the auto_wrap_policy is not provided, generate the default policy
     if auto_wrap_policy is None:
@@ -140,3 +140,41 @@ def prepare_fully_shard(
             if attr == 'verbose':
                 continue
             log.info(f'FSDP2: {attr}: {getattr(fsdp2_config, attr)}')
+
+
+def add_fsdp2_oom_hooks(model: nn.Module) -> tuple[list[torch.utils.hooks.RemovableHandle], dict[str, nn.Module]]:
+    """Add OOM hooks to the valid FSDP2-wrapped modules in a ComposerModel and return the named modules.
+
+    Args:
+        model (nn.Module): The model to add OOM hooks to.
+
+    Returns:
+        list[torch.utils.hooks.RemovableHandle]: A list of removable hook handles for the OOM hooks.
+        dict[str, nn.Module]: A dictionary of valid named modules in the ComposerModel.
+    """
+    assert isinstance(model, ComposerModel), f'{type(model)} is not a ComposerModel'
+    hook_handles = add_fsdp_oom_hooks(model, fsdp_config_version=2)
+
+    named_modules = {}
+    direct_children = get_direct_children_from_composer_model(model)
+    for child in direct_children:
+        named_modules.update(dict(child.named_modules()))
+    return hook_handles, named_modules
+
+
+def unset_fsdp2_state(model: nn.Module) -> None:
+    """Recover the FSDP2 state of the model to its original state.
+
+    This is needed because FSDP2 is stateful since they expect training to fail when we hit OOM.
+    Since with auto-microbatching, training does not fail, we need to set the FSDP2 state to `IDLE`
+    so it can continue transitioning between states correctly. Otherwise, there's a chance that it won't
+    transition correctly and either hang/raise unexpected errors.
+
+    Args:
+        model (nn.Module): The model to recover the FSDP2 state of.
+    """
+    for module in model.modules():
+        if isinstance(module, FSDPModule):
+            state = fully_shard.state(module)  # type: ignore
+            state._state_ctx.post_backward_final_callback_queued = False  # type: ignore
+            state._training_state = TrainingState.IDLE  # type: ignore

composer/distributed/prepare_distributed.py

@@ -12,7 +12,7 @@
 from torch.distributed.fsdp.wrap import CustomPolicy
 
 from composer.distributed.activation_checkpointing import apply_ac, generate_composer_model_check_fn
-from composer.distributed.fsdp2 import prepare_fully_shard
+from composer.distributed.fsdp2 import add_fsdp2_oom_hooks, prepare_fully_shard
 from composer.distributed.fsdp2_utils import generate_composer_model_policy, sync_optimizer_and_model_params
 from composer.distributed.param_init import meta_init
 from composer.models import ComposerModel
@@ -91,7 +91,8 @@ def parallelize_composer_model(
     composer_model: ComposerModel,
     optimizer: Optional[torch.optim.Optimizer],
     config: FSDP2Config,
-):
+    auto_microbatching: bool = False,
+) -> tuple[list, dict]:
     """Prepare a ComposerModel for distributed training.
 
     NOTE we apply parallelization to each of the composer model's submodules to provide compatibility with models defined for FSDP1.
@@ -104,6 +105,12 @@ def parallelize_composer_model(
         composer_model (ComposerModel): The ComposerModel to prepare for distributed training.
         optimizer (Optional[torch.optim.Optimizer]): The optimizer to use for distributed training.
         config (FSDP2Config): The configuration for distributed training. Currently only FSDP2Config is supported.
+        auto_microbatching (bool): Whether to use auto microbatching.
+
+    Returns:
+        tuple[list, dict]: A tuple containing:
+            - A list of removable hook handles for the OOM hooks if auto_microbatching is enabled
+            - A dictionary mapping module names to modules after fully sharding
     """
     assert isinstance(composer_model, ComposerModel), f'{type(composer_model)} is not a ComposerModel'
     activation_checkpointing_check_fn = generate_composer_model_check_fn(
@@ -117,3 +124,9 @@ def parallelize_composer_model(
         activation_checkpointing_check_fn=activation_checkpointing_check_fn,
         param_init_fn=meta_init,
     )
+
+    hook_handles = []
+    named_modules = {}
+    if auto_microbatching:
+        hook_handles, named_modules = add_fsdp2_oom_hooks(composer_model)
+    return hook_handles, named_modules

composer/distributed/shared_utils.py

@@ -0,0 +1,166 @@
+# Copyright 2022 MosaicML Composer authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared utilities for distributed training."""
+
+import functools
+from typing import Callable, Optional
+
+import torch
+from torch.distributed.fsdp import FSDPModule, FullyShardedDataParallel
+from torch.utils.hooks import RemovableHandle
+from torchmetrics import Metric, MetricCollection
+
+from composer.devices import Device
+from composer.models import ComposerModel
+from composer.utils import dist, get_device
+
+
+def get_valid_fsdp_module_types() -> dict[int, type]:
+    """Returns a dictionary of valid FSDP module types based on the torch version.
+
+    Returns:
+        dict: Dictionary of valid FSDP module types.
+    """
+    return {
+        1: FullyShardedDataParallel,
+        2: FSDPModule,
+    }
+
+
+def get_direct_children_from_composer_model(model: ComposerModel) -> list[torch.nn.Module]:
+    """Returns a list of valid direct children from a ComposerModel.
+
+    A valid direct child for a ComposerModel is a module that's not a Metric or MetricCollection.
+
+    Returns:
+        list: List of valid direct children from a ComposerModel.
+    """
+    assert isinstance(model, ComposerModel)
+    direct_children = []
+    for child in model.children():
+        if isinstance(child, (Metric, MetricCollection)):
+            continue
+        direct_children.append(child)
+
+    return direct_children
+
+
+def generate_oom_hook(device: Device) -> Callable:
+    """Generate a hook that checks if any other rank hit an OOM.
+
+    We 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(1/2) block, we wrap every
+    FSDP(1/2) block with a hook that checks if any other rank OOMed.
+
+    Here's an example of why this is needed using a simple 2-GPU setup and how it handles OOM issues during auto microbatching.
+    Note that the line numbers can be (slightly) off based on future changes made to the code.
+
+    - Rank 0: Layer 1 works fine
+    - Rank 1: Layer 1 works fine
+    - Rank 0: Layer 2 OOMs
+        - Rank 0 raises an error _is_cuda_oom() [[trainer.py:2756]]
+        - Rank 0 sets found_cuda_oom to 1 [[trainer.py:2758]]
+        - Rank 0 creates found_cuda_oom_tensor = [1] and calls all_reduce on it with reduce_operation='MAX' [[trainer.py:2773]]
+    - Rank 1: Layer 2 works fine until a hook handle is hit
+        - Rank 1 sets found_cuda_oom_tensor = [0] [[shared_utils.py:85]]
+        - Rank 1 calls all_reduce to set found_cuda_oom_tensor to max([0, 1]) = 1 [[fsdp2.py:73]]
+        - Rank 1 sees that found_cuda_oom == 1 [[shared_utils.py:87]]
+    - Rank 0:
+        - Rank 0 creates all_ranks_finished_tensor = [1] and calls all_reduce on it with reduce_operation='MIN' [[trainer.py:2780]]
+        - Rank 0 sees that all_ranks_finished == 0 (since rank 1 is still in mid-batch) [[trainer.py:2781]]
+        - Rank 0 continues in the (while not all_ranks_finished) loop [[trainer.py:2771]]
+    - Rank 1:
+        - Rank 1 creates all_ranks_finished_tensor = [0] and calls all_reduce on it with reduce_operation='MIN' [[shared_utils.py:89]]
+        - Rank 1 sees that all_ranks_finished == 0 (since this rank is still in the batch) [[shared_utils.py:90]]
+        - Rank 1 sees that found_cuda_oom == 1, so it raises an error saying that a different rank OOMed [[shared_utils.py:93]]
+    - Rank 0:
+        - In the next round of the while loop, found_cuda_oom_tensor = [1] and calls all_reduce on it with reduce_operation='MAX' [[trainer.py:2773]]
+    - Rank 1:
+        - Rank 1 sees the error that was raised earlier (OOM on other rank) and sets found_cuda_oom to 1 [[trainer.py:2755]]
+        - Rank 1 creates found_cuda_oom_tensor = [1] and calls all_reduce on it with reduce_operation='MAX' [[trainer.py:2773]]
+        - As expected, found_cuda_oom == 1 [[trainer.py:2776]]
+    - Rank 0:
+        - Rank 0 creates all_ranks_finished_tensor = [1] (since it's in the same while loop as before) and calls all_reduce on it with reduce_operation='MIN' [[trainer.py:2780]]
+        - Rank 0 sees that all_ranks_finished = 1 (as we are in the same part of the trainer code as Rank 1, Rank 1 returns the same value) [[trainer.py:2782]]
+        - Rank 0 exits the while loop and adjusts the device_train_microbatch_size to half of the previous value [[trainer.py:2790]]
+    - Rank 1:
+        - Rank 1 creates all_ranks_finished_tensor = [1] (since it's finished the batch with an error) and calls all_reduce on it with reduce_operation='MIN' [[trainer.py:2780]]
+        - Rank 1 sees that all_ranks_finished == 1 (since this rank is finished the batch) [[trainer.py:2781]]
+        - Rank 1 exits the while loop and adjusts the device_train_microbatch_size to half of the previous value [[trainer.py:2790]]
+
+    Args:
+        device (torch.device): The device to check for OOM.
+
+    Returns:
+        Callable: The hook that checks if any other rank hit an OOM.
+    """
+
+    def sync_hook(*args, device: Device):
+        # Check if any other rank hit an OOM
+        found_cuda_oom_tensor = 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 = 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('CUDA out of memory encountered on a different rank')
+
+    return functools.partial(sync_hook, device=device)
+
+
+def add_fsdp_oom_hooks(model, fsdp_config_version: int, device: Optional[Device] = None) -> list[RemovableHandle]:
+    """Add OOM hooks to the model and return the list of handles.
+
+    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`
+
+    View https://github.com/mosaicml/composer/pull/3510 for more details.
+
+    Args:
+        model (torch.nn.Module): The model to add the hooks to. This can be a ComposerModel and in that scenario, we need to add hooks to valid children.
+        fsdp_config_version (int): The version of the FSDP config to use. This should be either 1 or 2.
+        device (torch.device): The device that the module is on. If None, the current rank's device will be used.
+
+    Returns:
+        list[RemovableHandle]: The list of RemovableHandles for the hooks.
+    """
+    hook_handles = []
+    if device is None:
+        device = get_device()
+    hook = generate_oom_hook(device)
+
+    # Gets the valid children if the input is a ComposerModel
+    root_modules_for_hooks = []
+    if isinstance(model, ComposerModel):
+        root_modules_for_hooks = get_direct_children_from_composer_model(model)
+    else:
+        root_modules_for_hooks.append(model)
+
+    # In FSDP2, we don't support backward_prefetch=BACKWARD_POST, so we only need to add the OOM hooks to the
+    # forward pre and backward pre hooks.
+    # TODO: In FSDP1, we might not need the non-FSDP wrapped backward hook either, but we'll keep it for now until further investigation.
+    # TODO: If we want to reduce as many potential deadlocks as possible, we may need to add hooks before all blocking collectives:
+    #   - register_forward_pre_hook (before blocking all_gather)
+    #   - register_full_backward_pre_hook (before blocking all_gather)
+    #   - register_full_backward_hook (before blocking reduce_scatter)
+    # In all of these cases, some combination of no activation checkpointing/offloading, reshard_after_forward=False, or high gradient memory cost
+    # could result in edge-case OOMs and deadlocks.
+    fsdp_module_type = get_valid_fsdp_module_types()[fsdp_config_version]
+    for root_module in root_modules_for_hooks:
+        for module in root_module.modules():
+            if isinstance(module, fsdp_module_type):
+                hook_handles.append(module.register_forward_pre_hook(hook, prepend=True))  # type: ignore
+                hook_handles.append(module.register_full_backward_pre_hook(hook, prepend=True))  # type: ignore
+            elif fsdp_config_version == 1:
+                hook_handles.append(module.register_full_backward_hook(hook))  # type: ignore
+
+    return hook_handles