harrisonvanderbyl / rwkvstic

Framework agnostic python runtime for RWKV models
https://hazzzardous-rwkv-instruct.hf.space
MIT License
144 stars 18 forks source link

TORCH_STREAM not working #15

Closed MichelNivard closed 1 year ago

MichelNivard commented 1 year ago

I cant use TORCH_STREAM, but I can select it via inquirer I think (but thats a very uneducated guess) its because TORCH_STREAM is not defined in the init.py in /agnostic/backends:

Backends: Dict[str, module] = {
    "tensorflow(cpu/gpu)": RWKVTFOps,
    "pytorch(cpu/gpu)": RWKVCudaOps,
    "numpy(cpu)": RWKVNumpyOps,
    "jax(cpu/gpu/tpu)": RWKVJaxOps,
    "pytorch-deepspeed(gpu)": RWKVCudaDeepspeedOps,
    "pytorch-quant(gpu-8bit)": RWKVCudaQuantOps,
    "pytorch-stream(gpu-config-vram)": RWKVStreamBigOps,
    "pytorch-split(2xgpu)": RWKVSplitCudaOps,
    "export-torchscript": RWKVPTTSExportOps,
    "export-tensorflow": RWKVTFExport,
    "onnx_export": RWKVOnnxOps,
    # "coreml": RWKVCoreMLOps
}

TORCH = "pytorch(cpu/gpu)"
JAX = "jax(cpu/gpu/tpu)"
TF = "tensorflow(cpu/gpu)"
TORCH_DEEPSPEED = "pytorch-deepspeed(gpu)"
TORCH_QUANT = "pytorch-quant(gpu-8bit)"
TORCH_SPLIT = "pytorch-split(2xgpu)"
TORCH_EXPORT = "export-torchscript"
TF_EXPORT = "export-tensorflow"
ONNX_EXPORT = "onnx_export"
NUMPY = "numpy(cpu)"

# COREML = "coreml"