Function wrapper

fastdeploy/model_executor/load_weight_utils.py:173–214

(*args, **kwargs)

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            sig = inspect.signature(func)
            bound_args = sig.bind(*args, **kwargs)
            bound_args.apply_defaults()
            fd_config = bound_args.arguments.get(config_arg_name, None)
            model = bound_args.arguments.get(model_arg_name, None)
            enable_cache, weight_cache_dir, _ = is_weight_cache_enabled(fd_config)
            assert fd_config is not None, "fd_config cannot be None"
            assert model is not None, "model cannot be None"
            if enable_cache:
                tp_weight_cache_dir = os.path.join(
                    weight_cache_dir, f"rank{str(fd_config.parallel_config.tensor_parallel_rank)}"
                )
                context = multi_switch_config_context((fd_config.model_config, "model", tp_weight_cache_dir))
            else:
                context = contextlib.nullcontext()

            with context:
                result = func(*args, **kwargs)

            if envs.FD_ENABLE_MODEL_LOAD_CACHE:
                if not (
                    fd_config.quant_config is not None and getattr(fd_config.quant_config, "is_checkpoint_bf16", False)
                ):
                    # Save cache only for dynamic quantization
                    return result
                if weight_cache_dir is None:
                    return result
                tp_weight_cache_dir = os.path.join(
                    weight_cache_dir, f"rank{str(fd_config.parallel_config.tensor_parallel_rank)}"
                )
                if not os.path.exists(tp_weight_cache_dir):
                    logger.info(f"Saving model to {tp_weight_cache_dir}")
                    os.makedirs(
                        tp_weight_cache_dir,
                        exist_ok=True,
                    )
                    _save_model(model.state_dict(), os.path.join(tp_weight_cache_dir, "cache.pdparams"))
                else:
                    reason = "weights already cached" if envs.FD_ENABLE_MODEL_LOAD_CACHE else "cache disabled"
                    logger.info(f"Skip saving ,{reason}")
            return result

        return wrapper