diff --git a/composer/distributed/dist_strategy.py b/composer/distributed/dist_strategy.py index ae621491f7..2d9fa152a0 100644 --- a/composer/distributed/dist_strategy.py +++ b/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'): diff --git a/composer/distributed/fsdp2.py b/composer/distributed/fsdp2.py index 7bb8c9617e..a96cf3cd37 100644 --- a/composer/distributed/fsdp2.py +++ b/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 diff --git a/composer/distributed/prepare_distributed.py b/composer/distributed/prepare_distributed.py index cd6af25627..45ba525edd 100644 --- a/composer/distributed/prepare_distributed.py +++ b/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 diff --git a/composer/distributed/shared_utils.py b/composer/distributed/shared_utils.py new file mode 100644 index 0000000000..675f753731 --- /dev/null +++ b/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 diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py index c0154d3fed..b52cce4c88 100644 --- a/composer/trainer/trainer.py +++ b/composer/trainer/trainer.py @@ -79,6 +79,7 @@ prepare_fsdp_module, prepare_tp_module, ) +from composer.distributed.shared_utils import generate_oom_hook, get_valid_fsdp_module_types from composer.loggers import ( ConsoleLogger, Logger, @@ -427,44 +428,19 @@ def _update_num_consecutive_thrashes(state: State, num_consecutive_thrashes: int 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): +def _readd_fsdp_sync_hooks(fsdp_modules: dict[str, torch.nn.Module], fsdp_config_version: int, 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) + fsdp_module_type = get_valid_fsdp_module_types()[fsdp_config_version] for module in fsdp_modules.values(): - if isinstance(module, FullyShardedDataParallel): + if isinstance(module, fsdp_module_type): 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: + elif fsdp_config_version == 1: automicrobatch_fsdp_hook_handles.append(module.register_full_backward_hook(sync_hook)) return automicrobatch_fsdp_hook_handles @@ -1859,10 +1835,11 @@ def _wrap_model_for_distributed( self.state.seed, ) case 2: - parallelize_composer_model( + self.state.automicrobatch_fsdp_hook_handles, self.state.fsdp_modules = parallelize_composer_model( model, optimizers, self.state.fsdp_config, # type: ignore + auto_microbatching, ) case _: raise ValueError(f'Unsupported FSDP config version: {self.state.fsdp_config_version}') @@ -2713,7 +2690,7 @@ def _train_batch(self, use_grad_scaling: bool) -> dict[str, torch.Tensor]: device_batch = self.state.batch # Define sync hook for FSDP modules if automicrobatching is on - sync_hook = _create_sync_hook(self.state) + sync_hook = generate_oom_hook(self.state.device) original_microbatch_size = self.state.device_train_microbatch_size oom_found_this_batch = False @@ -2790,9 +2767,12 @@ def _train_batch(self, use_grad_scaling: bool) -> dict[str, torch.Tensor]: 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, + fsdp_config_version=self.state.fsdp_config_version, + sync_hook=sync_hook, ) _adjust_device_train_microbatch_size(self.state) + from composer.distributed.fsdp2 import unset_fsdp2_state + unset_fsdp2_state(self.state.model) self.num_consecutive_thrashes = 0 self.num_consecutive_non_OOM_batches = 0 oom_found_this_batch = True @@ -2810,7 +2790,8 @@ def _train_batch(self, use_grad_scaling: bool) -> dict[str, torch.Tensor]: 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, + fsdp_config_version=self.state.fsdp_config_version, + sync_hook=sync_hook, ) _adjust_device_train_microbatch_size(self.state) self.num_consecutive_thrashes = 0 @@ -3619,9 +3600,13 @@ def _eval_loop( # If training occurs after evaluation, readd hooks in case of memory spike if self.state.auto_microbatching: - sync_hook = _create_sync_hook(self.state) + sync_hook = generate_oom_hook(self.state.device) 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.state.automicrobatch_fsdp_hook_handles = _readd_fsdp_sync_hooks( + self.state.fsdp_modules, + fsdp_config_version=self.state.fsdp_config_version, + sync_hook=sync_hook, + ) self.num_consecutive_non_OOM_batches = 0 def _use_grad_scaling(self, precision: Union[str, Precision], scaler: Optional[GradScaler]) -> bool: diff --git a/tests/common/__init__.py b/tests/common/__init__.py index 1c6ad74545..0952dcf164 100644 --- a/tests/common/__init__.py +++ b/tests/common/__init__.py @@ -24,6 +24,7 @@ EmbeddedWeightTiedModel, EmptyModel, EvenSimplerMLP, + OOMComposerClassifier, PartialWeightTiedModel, SimpleComposerMLP, SimpleConvModel, @@ -79,4 +80,5 @@ def get_module_subclasses(module: types.ModuleType, cls: type) -> list[type]: 'TPSimpleComposerMLP', 'ComposerCounterModel', 'CountModule', + 'OOMComposerClassifier', ] diff --git a/tests/common/models.py b/tests/common/models.py index 99fec7350f..89546e140d 100644 --- a/tests/common/models.py +++ b/tests/common/models.py @@ -16,6 +16,7 @@ from composer.metrics import CrossEntropy, MIoU from composer.metrics.nlp import LanguageCrossEntropy, MaskedAccuracy from composer.models import ComposerClassifier, HuggingFaceModel, Initializer +from composer.utils import dist if TYPE_CHECKING: from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast @@ -180,6 +181,40 @@ def __init__( self.module = module +class OOMComposerClassifier(ComposerClassifier): + """A model that will raise an OOM error on rank 1 when forward is called in expected situations. + + This is used to test the auto microbatching code and will always fail on rank 1 if the microbatch size + is greater than the viable microbatch size. + """ + + def __init__( + self, + num_layers: int, + num_classes: int, + device: Union[str, torch.device], + always_fail: bool = False, + viable_microbatch_size: int = 32, + ): + module = torch.nn.Sequential( + *[torch.nn.Linear(num_classes, num_classes, device=device) for _ in range(num_layers)], + ) + super().__init__( + num_classes=num_classes, + module=module, + ) + self.module = module + self.always_fail = always_fail + self.viable_microbatch_size = viable_microbatch_size + + def forward(self, batch: tuple[torch.Tensor, Any]) -> torch.Tensor: + inputs, _ = batch + outputs = self.module(inputs) + if dist.get_global_rank() == 1 and (self.always_fail or inputs.shape[0] > self.viable_microbatch_size): + raise RuntimeError('CUDA out of memory') + return outputs + + # Like SimpleComposerMLP but saves each layer which is necessary to TP to it. class TPSimpleComposerMLP(ComposerClassifier): diff --git a/tests/trainer/test_fsdp2.py b/tests/trainer/test_fsdp2.py index 1f847437dc..81fe8fb31d 100644 --- a/tests/trainer/test_fsdp2.py +++ b/tests/trainer/test_fsdp2.py @@ -7,13 +7,16 @@ import pytest import torch from torch.distributed._tensor import DTensor +from torch.distributed.fsdp import FSDPModule from torch.utils.data import DataLoader +from torch.utils.hooks import RemovableHandle from composer.models import ComposerClassifier from composer.trainer.trainer import Trainer from composer.utils import dist, load_checkpoint from composer.utils.parallelism import FSDP2Config, FSDPConfig, ParallelismConfig from tests.common import ( + OOMComposerClassifier, PartialWeightTiedModel, RandomClassificationDataset, SimpleComposerMLP, @@ -27,15 +30,21 @@ def create_trainer_with_model( model: ComposerClassifier, num_classes: int = 10, + dataset_size: int = 2, max_duration: str = '10ep', use_fsdp2: bool = True, optimizer: Optional[torch.optim.Optimizer] = None, activation_checkpointing: bool = False, activation_cpu_offload: bool = False, + auto_microbatching: bool = False, ) -> Trainer: """Helper function to create a Trainer with a model, dataloader, and FSDP2 configuration.""" - dataset = RandomClassificationDataset(shape=(num_classes,), size=2, num_classes=num_classes) - dataloader = DataLoader(dataset, sampler=dist.get_sampler(dataset)) + dataset = RandomClassificationDataset(shape=(num_classes,), size=dataset_size, num_classes=num_classes) + dataloader = DataLoader( + dataset, + sampler=dist.get_sampler(dataset), + batch_size=dataset_size // 2, + ) # use 2 batches per epoch parallelism_config = ParallelismConfig() if use_fsdp2: @@ -53,7 +62,9 @@ def create_trainer_with_model( train_dataloader=dataloader, max_duration=max_duration, parallelism_config=parallelism_config, + device_train_microbatch_size='auto' if auto_microbatching else None, ) + return trainer @@ -303,3 +314,108 @@ def test_fsdp2_optimizer_raises_error_when_optimizer_modules_dont_match( # We check with `optimizer.param_id.` (with the period) since `optimizer.param_id` exists # by default in the error message's legend assert 'optimizer.param_id.' in str(e.value) + + +@world_size(2) +@pytest.mark.gpu +@pytest.mark.filterwarnings("ignore:`device_train_microbatch_size='auto'` may potentially fail with unexpected.*") +@pytest.mark.parametrize( + 'use_alternate,num_layers,expected_num_hooks', + [ + # 3 children modules wrapped * 2 new hook handles per module + (False, 3, 3 * 2), + # 2 children modules wrapped * 2 new hook handles per module + (True, 3, 2 * 2), + ], +) +def test_fsdp2_has_right_number_of_hooks( + world_size: int, + use_alternate: bool, + num_layers: int, + expected_num_hooks: int, +): + """Test FSDP2 has the right number of hooks.""" + del world_size + + num_classes = 10 + model = OOMComposerClassifier(num_layers, num_classes, device='cuda') + + # Wrap the module as we expect + model.module._fsdp_wrap = False # type: ignore + for i, child in enumerate(model.module.children()): + if use_alternate: + child._fsdp_wrap = True if i % 2 == 0 else False # type: ignore + else: + child._fsdp_wrap = True # type: ignore + + # Assert that the number of hooks returned is correct + trainer = create_trainer_with_model( + model=model, + num_classes=num_classes, + use_fsdp2=True, + auto_microbatching=True, + ) + hook_handles = trainer.state.automicrobatch_fsdp_hook_handles + error_msg = 'Expected {} OOM hooks, but got {}' + assert len(hook_handles) == expected_num_hooks, error_msg.format(expected_num_hooks, len(hook_handles)) + + # Assert that all hook handles are RemovableHandle + for hook_handle in hook_handles: + assert isinstance(hook_handle, RemovableHandle), f'Expected RemovableHandle, but got {type(hook_handle)}' + + # Assert number of hooks on each module + # Note: reshard_after_forward doesn't change the number of backward_hooks, it just changes the existing hooks do so the numbers + # below are the same for both reshard_after_forward = True and False. + error_msg = 'Expected {} forward pre hooks on module {}, but got {}' + for child in model.module.children(): + if isinstance(child, FSDPModule): # type: ignore + # Hooks exist for FSDP wrapped modules + assert len(child._forward_pre_hooks) == 2, error_msg.format(2, child, len(child._forward_pre_hooks)) + assert len(child._backward_pre_hooks) == 1, error_msg.format(1, child, len(child._backward_pre_hooks)) + assert len(child._backward_hooks) == 0, error_msg.format(0, child, len(child._backward_hooks)) + else: + # No hooks on non-FSDP wrapped modules + assert len(child._forward_pre_hooks) == 0, error_msg.format(0, child, len(child._forward_pre_hooks)) + assert len(child._backward_pre_hooks) == 0, error_msg.format(0, child, len(child._backward_pre_hooks)) + assert len(child._backward_hooks) == 0, error_msg.format(0, child, len(child._backward_hooks)) + + +@world_size(2) +@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*') +def test_fsdp2_auto_microbatching_handles_cuda_failures( + world_size: int, +): + """Test FSDP2 auto-microbatching handles CUDA OOM failures.""" + del world_size + + # This will always fail (rank 1 will fail on all batch sizes) + num_classes = 10 + model = OOMComposerClassifier(3, num_classes, device='cuda', always_fail=True) + for child in model.module.children(): + child._fsdp_wrap = True # type: ignore + trainer = create_trainer_with_model( + model=model, + num_classes=num_classes, + use_fsdp2=True, + auto_microbatching=True, + dataset_size=256, + max_duration='1ba', + ) + with pytest.raises(RuntimeError, match='.*The train loop failed with an internal microbatch of size 1.*'): + trainer.fit() + + # This will succeed as rank 1 will not fail if the microbatch size <= 32 + model = OOMComposerClassifier(3, num_classes, device='cuda', always_fail=False, viable_microbatch_size=32) + for child in model.module.children(): + child._fsdp_wrap = True # type: ignore + trainer = create_trainer_with_model( + model=model, + num_classes=num_classes, + use_fsdp2=True, + auto_microbatching=True, + dataset_size=256, + max_duration='1ba', + ) + trainer.fit()