def _get_granularity_recursively(self, module):
        """This function is used to recursively obtain the granularity of each module."""

        # avoid setting as leaf for particularly large models, even if the granularity is very small
        # an oversized leaf module increases the number of live parameters, introducing memory overhead
        Z3_MAX_LEAF_SIZE = 1e9

        if not list(module.parameters()):
            # skip Modules without parameters, such as GELU, etc.
            module.ds_model_granularity = sys.maxsize
            return 0, 0

        num_layers = 0
        num_params = 0
        num_params += sum(p.ds_numel for p in module.parameters(recurse=False))
        if not any(module.children()):
            # torch leaf module
            module.ds_model_granularity = sys.maxsize
            return 1, num_params

        for child in module.children():
            layers_in_child, params_in_child = self._get_granularity_recursively(child)
            num_layers += layers_in_child
            num_params += params_in_child

        if module.__class__.__name__ in torch.nn.modules.container.__all__:
            # Do not set container modules like ModuleList as leaf modules
            # as this will prevent hooks from being set on their children
            # and they may do not invoke the forward method
            module.ds_model_granularity = sys.maxsize
            return num_layers, num_params

        num_layers += 1
        ds_model_granularity = (num_params // num_layers) if num_params <= Z3_MAX_LEAF_SIZE else sys.maxsize
        module.ds_model_granularity = ds_model_granularity
        # module.ds_model_num_layers = num_layers
        # module.ds_model_num_params = num_params
        if self.min_granularity_value > ds_model_granularity:
            self.min_granularity_value = ds_model_granularity
            self.min_granularity_layer = module.__class__.__name__
        self.granularity_info.add(f"{module.__class__.__name__.ljust(30)}|{str(ds_model_granularity).rjust(20)}")

        return num_layers, num_params
This function is used to recursively obtain the granularity of each module.
docker/parameter_offload.py