Function parallelize_model

composer/distributed/prepare_distributed.py:138–195

Prepare a model for distributed training. This function currently applies FSDP2 to the model, initializes parameters, and optionally applies activation checkpointing. It handles parameter consistency between the optimizer and model when an optimizer is provided. Args: model

(
    model: torch.nn.Module,
    config: FSDP2Config,
    optimizer: Optional[torch.optim.Optimizer] = None,
    precision: Optional[Union[str, Precision]] = None,
    fsdp_wrap_policy: Optional[CustomPolicy] = None,
    activation_checkpointing_check_fn: Optional[Callable] = None,
    param_init_fn: Callable[[torch.nn.Module], None] = lambda m: None,
)

def parallelize_model(
    model: torch.nn.Module,
    config: FSDP2Config,
    optimizer: Optional[torch.optim.Optimizer] = None,
    precision: Optional[Union[str, Precision]] = None,
    fsdp_wrap_policy: Optional[CustomPolicy] = None,
    activation_checkpointing_check_fn: Optional[Callable] = None,
    param_init_fn: Callable[[torch.nn.Module], None] = lambda m: None,
):
    """Prepare a model for distributed training.

    This function currently applies FSDP2 to the model, initializes parameters,
    and optionally applies activation checkpointing. It handles parameter consistency between the optimizer
    and model when an optimizer is provided.

    Args:
        model (torch.nn.Module): The model to prepare for distributed training.
        config (FSDP2Config): The configuration for FSDP distributed training.
        optimizer (Optional[torch.optim.Optimizer]): The optimizer to synchronize with the model.
            If provided, parameter states will be properly synchronized during sharding.
        precision (Precision): The precision to use for the model. Defaults to AMP_FP16 for GPU and FP32 for CPU.
            It doesn't have an optional type because `parallelize_composer_model` already sets the precision.
        fsdp_wrap_policy (Optional[CustomPolicy]): Custom policy to determine which modules should
            be wrapped with FSDP. If None, default wrapping behavior is used.
        activation_checkpointing_check_fn (Optional[Callable]): Function that determines whether a
            module's activations should be checkpointed or offloaded. Only used when activation
            checkpointing or CPU offloading is enabled in the config.
        param_init_fn (Callable[[torch.nn.Module], None]): Function to initialize model parameters
            after FSDP wrapping. Defaults to a no-op function.

    Raises:
        ValueError: If the config is not an FSDP2Config or if activation_checkpointing_check_fn is provided
            but neither activation_checkpointing nor activation_cpu_offload is enabled in the config.
    """
    if precision is None:
        precision = Precision.AMP_FP16 if isinstance(get_device(), DeviceGPU) else Precision.FP32
    elif isinstance(precision, str):
        precision = Precision(precision)
    _validate_precision(precision, get_device())

    if activation_checkpointing_check_fn is not None:
        if not config.activation_checkpointing and not config.activation_cpu_offload:
            raise ValueError(
                'Activation checkpointing or offloading must be enabled if activation_checkpointing_check_fn is provided',
            )

    if config.activation_checkpointing or config.activation_cpu_offload:
        apply_ac(
            model,
            config.activation_checkpointing,
            config.activation_cpu_offload,
            activation_checkpointing_check_fn,
        )

    # Use the context manager for optimizer synchronization if optimizer is provided
    with sync_optimizer_and_model_params(optimizer, model) if optimizer is not None else nullcontext():
        _parallelize_model_helper(model, config, precision, fsdp_wrap_policy, param_init_fn)
        # NOTE appy_ac can not be included in this context as it would wrap and replace the sub-modules thus disqualify FQN of params