def _set_leaf_by_threshold_preorder(self, module, granularity_treshhold):
        '''Set modules as leaf modules based on the threshold, prioritizing parent nodes.'''

        num_params = sum(p.ds_numel for p in module.parameters())
        if num_params == 0:
            # skip Modules without parameters, such as GELU, etc.
            return
        if module.ds_model_granularity <= granularity_treshhold:
            set_z3_leaf_module(module, True)
            self.z3_leaf_layers.append(module)
            return

        for sub_module in module.children():
            self._set_leaf_by_threshold_preorder(sub_module, granularity_treshhold)
Set modules as leaf modules based on the threshold, prioritizing parent nodes.
docker/parameter_offload.py