Function _get_expected_safetensors_files

tests/testing_utils.py:1641–1722

(
        module: torch.nn.Module,
        offload_to_disk_path: str,
        offload_type: str,
        num_blocks_per_group: int | None = None,
        block_modules: list[str] | None = None,
        module_prefix: str = "",
    )

    )

    def _get_expected_safetensors_files(
        module: torch.nn.Module,
        offload_to_disk_path: str,
        offload_type: str,
        num_blocks_per_group: int | None = None,
        block_modules: list[str] | None = None,
        module_prefix: str = "",
    ) -> set[str]:
        expected_files = set()

        def get_hashed_filename(group_id: str) -> str:
            short_hash = _compute_group_hash(group_id)
            return os.path.join(offload_to_disk_path, f"group_{short_hash}.safetensors")

        if offload_type == "block_level":
            if num_blocks_per_group is None:
                raise ValueError("num_blocks_per_group must be provided for 'block_level' offloading.")

            block_modules_set = set(block_modules) if block_modules is not None else set()

            modules_with_group_offloading = set()
            unmatched_modules = []
            for name, submodule in module.named_children():
                if name in block_modules_set:
                    new_prefix = f"{module_prefix}{name}." if module_prefix else f"{name}."
                    submodule_files = _get_expected_safetensors_files(
                        submodule, offload_to_disk_path, offload_type, num_blocks_per_group, block_modules, new_prefix
                    )
                    expected_files.update(submodule_files)
                    modules_with_group_offloading.add(name)

                elif isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
                    for i in range(0, len(submodule), num_blocks_per_group):
                        current_modules = submodule[i : i + num_blocks_per_group]
                        if not current_modules:
                            continue
                        group_id = f"{module_prefix}{name}_{i}_{i + len(current_modules) - 1}"
                        expected_files.add(get_hashed_filename(group_id))
                        for j in range(i, i + len(current_modules)):
                            modules_with_group_offloading.add(f"{name}.{j}")
                else:
                    unmatched_modules.append(submodule)

            parameters = _gather_parameters_with_no_group_offloading_parent(module, modules_with_group_offloading)
            buffers = _gather_buffers_with_no_group_offloading_parent(module, modules_with_group_offloading)

            if len(unmatched_modules) > 0 or len(parameters) > 0 or len(buffers) > 0:
                expected_files.add(get_hashed_filename(f"{module_prefix}{module.__class__.__name__}_unmatched_group"))

        elif offload_type == "leaf_level":
            # Handle leaf-level module groups
            for name, submodule in module.named_modules():
                if isinstance(submodule, _GO_LC_SUPPORTED_PYTORCH_LAYERS):
                    # These groups will always have parameters, so a file is expected
                    expected_files.add(get_hashed_filename(name))

            # Handle groups for non-leaf parameters/buffers
            modules_with_group_offloading = {