wehub-resource-sync
ec436095dd
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
30 lines
932 B
Python
30 lines
932 B
Python
from functools import wraps
|
|
import torch, yaml, pathlib
|
|
|
|
import os, sys
|
|
project_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
|
sys.path.insert(0, project_dir)
|
|
|
|
from ktransformers.util.globals import GLOBAL_CONFIG
|
|
|
|
# print(f"start_sit: {GLOBAL_CONFIG._config['mod']}")
|
|
|
|
def maybe_no_grad(_func=None):
|
|
# print(f"maybe_sit: {GLOBAL_CONFIG._config['mod']}")
|
|
|
|
def decorator(func):
|
|
# print(f"decorate_sit: {GLOBAL_CONFIG._config['mod']}")
|
|
def wrapper(*args, **kwargs):
|
|
# print(f"wrap_sit: {GLOBAL_CONFIG._config['mod']}")
|
|
if GLOBAL_CONFIG._config["mod"] == "sft":
|
|
return func(*args, **kwargs)
|
|
elif GLOBAL_CONFIG._config["mod"] == "infer":
|
|
with torch.no_grad():
|
|
return func(*args, **kwargs)
|
|
return wrapper
|
|
|
|
if _func is None:
|
|
return decorator
|
|
else:
|
|
return decorator(_func)
|