def dispatch_model(
    model,
    device_map: Union[str, dict],
    max_memory: Optional[dict] = None,
    no_split_module_classes: Optional[List[str]] = None,
    offload_folder: str = None,
    offload_buffers: bool = False,
    preload_module_classes: Optional[List[str]] = None,
):
    """Split and dispatch a model across devices.

    The function depends on the `accelerate` package. Refers to
    https://huggingface.co/docs/accelerate/main/en/usage_guides/big_modeling

    Args:
        model (torch.nn.Module): The model to dispatch.
        device_map (str | dict | None): A map that specifies where each
            submodule should go. It doesn't need to be refined to each
            parameter/buffer name, once a given module name is inside, every
            submodule of it will be sent to the same device. You can use
            `device_map="auto"` to automatically generate the device map.
            Defaults to None.
        max_memory (dict | None): A dictionary device identifier to maximum
            memory. Will default to the maximum memory available for each GPU
            and the available CPU RAM if unset. Defaults to None.
        no_split_module_classes (List[str] | None): A list of layer class names
            that should never be split across device (for instance any layer
            that has a residual connection). If None, try to get the settings
            from the model class. Defaults to None.
        offload_folder (str | None): If the `device_map` contains any value
            `"disk"`, the folder where we will offload weights.
        offload_buffers (bool): In the layers that are offloaded on the CPU
            or the hard drive, whether or not to offload the buffers as
            well as the parameters. Defaults to False.
        preload_module_classes (List[str] | None): A list of classes whose
            instances should load all their weights (even in the submodules) at
            the beginning of the forward. This should only be used for classes
            that have submodules which are registered but not called directly
            during the forward, for instance if a `dense` linear layer is
            registered, but at forward, `dense.weight` and `dense.bias` are
            used in some operations instead of calling `dense` directly.
            Defaults to None.
    """
    from accelerate import dispatch_model, infer_auto_device_map

    # Check valid device_map string.
    valid_map_option = ['auto', 'balanced', 'balanced_low_0', 'sequential']
    if isinstance(device_map, str) and device_map not in valid_map_option:
        raise ValueError('If passing a string for `device_map`, please choose '
                         f'from {valid_map_option}.')

    # Generate device map automatically
    if isinstance(device_map, str):
        if no_split_module_classes is None:
            no_split_module_classes = getattr(model, '_no_split_modules', None)
        if no_split_module_classes is None:
            raise ValueError(f'{model.__class__.__name__} does not support '
                             f"`device_map='{device_map}'` yet.")

        if device_map != 'sequential':
            from accelerate.utils import get_balanced_memory
            max_memory = get_balanced_memory(
                model,
                max_memory=max_memory,
                no_split_module_classes=no_split_module_classes,
                dtype=None,
                low_zero=(device_map == 'balanced_low_0'),
            )
            max_memory[0] *= 0.9
        device_map = infer_auto_device_map(
            model,
            max_memory=max_memory,
            no_split_module_classes=no_split_module_classes,
            dtype=None,
        )

    if 'disk' in device_map.values():
        if offload_folder is None:
            raise ValueError(
                'The current `device_map` had weights offloaded to the disk. '
                'Please provide an `offload_folder` for them.')
        os.makedirs(offload_folder, exist_ok=True)

    main_device = next(
        (d for d in device_map.values() if d not in ['cpu', 'disk']), 'cpu')

    model = dispatch_model(
        model,
        device_map=device_map,
        main_device=main_device,
        offload_dir=offload_folder,
        offload_buffers=offload_buffers,
        preload_module_classes=preload_module_classes,
    )
    if hasattr(model, 'data_preprocessor'):
        model.data_preprocessor._device = torch.device(main_device)
    return model
Split and dispatch a model across devices.

    The function depends on the `accelerate` package. Refers to
    https://huggingface.co/docs/accelerate/main/en/usage_guides/big_modeling

    Args:
        model (torch.nn.Module): The model to dispatch.
        device_map (str | dict | None): A map that specifies where each
            submodule should go. It doesn't need to be refined to each
            parameter/buffer name, once a given module name is inside, every
            submodule of it will be sent to the same device. You can use
            `device_map="auto"` to automatically generate the device map.
            Defaults to None.
        max_memory (dict | None): A dictionary device identifier to maximum
            memory. Will default to the maximum memory available for each GPU
            and the available CPU RAM if unset. Defaults to None.
        no_split_module_classes (List[str] | None): A list of layer class names
            that should never be split across device (for instance any layer
            that has a residual connection). If None, try to get the settings
            from the model class. Defaults to None.
        offload_folder (str | None): If the `device_map` contains any value
            `"disk"`, the folder where we will offload weights.
        offload_buffers (bool): In the layers that are offloaded on the CPU
            or the hard drive, whether or not to offload the buffers as
            well as the parameters. Defaults to False.
        preload_module_classes (List[str] | None): A list of classes whose
            instances should load all their weights (even in the submodules) at
            the beginning of the forward. This should only be used for classes
            that have submodules which are registered but not called directly
            during the forward, for instance if a `dense` linear layer is
            registered, but at forward, `dense.weight` and `dense.bias` are
            used in some operations instead of calling `dense` directly.
            Defaults to None.
    
mmpretrain/apis/utils.py