Add int4/int8 vicuna

[jinchen62]

• move brevitas_matmul_rhs_group_quant_library and vicuna pipeline to main file to avoid a decompose error

error: "<eval_with_key>.4":78:23: found an op that was marked as backend illegal
note: "<eval_with_key>.4":78:23: see current operation: %128 = "torch.aten._softmax"(%127, %40, %42) : (!torch.vtensor<[1,32,137,137],unk>, !torch.int, !torch.bool) -> !torch.vtensor<[1,32,137,137],unk>
note: "<eval_with_key>.4":78:23: this is likely due to DecomposeComplexOps being unable to decompose this op

• keep the format of brevitas〇matmul_rhs_group_quant to avoid a signature mismatch error

raise ValueError(f"Signature mismatch for {f.name!r}: expected {expected_signature!r}, got {signature!r}")
ValueError: Signature mismatch for 'brevitas〇matmul_rhs_group_quant〡shape': expected 'def brevitas〇matmul_rhs_group_quant〡shape(lhs: List[int], rhs: List[int], rhs_scale: List[int], rhs_zero_point: List[int], rhs_bit_width: int, rhs_group_size: int) -> List[int]:', got 'def brevitas〇matmul_rhs_group_quant〡shape('

• build torch-mlir from source with https://github.com/jinchen62/torch-mlir/tree/brevitas

TODO:

• check the canonicalization pass to make sure the tensors are processed properly
• apply canonicalization for int4 to unpack the tensors

[Abhishek-Varma]

Please update your patch as per the comment - precision in StableDiffusion's API won't be needed because it's already there.

[Abhishek-Varma]

ModuleNotFoundError: No module named 'brevitas_examples' - @jinchen62

Can you please resolve such failures we see in CI?

[powderluv]

We can add Brevitas to the requirements.txt with a git install

[Abhishek-Varma]

Please address one comment - rest looks good.

[Abhishek-Varma]

This doesn't seem correct. Can you confirm?

[jinchen62]

It does ignore this file from the format check. How does it look not corret for you?

[powderluv]

if you add it here it wont be detected by git. You want to add it to the black format command line with a --exclude or something

[Abhishek-Varma]

Please address the following sub-comments and also make changes to the following file in order to integrate with WebUI as well :-

Replace the set of lines at : https://github.com/nod-ai/SHARK/blob/91ab594744ffbe982ababe340fcd208923ecae48/apps/stable_diffusion/web/ui/stablelm_ui.py#L45-L47

With :

from apps.language_models.scripts.vicuna import (
            UnshardedVicuna,
)            

Replace the following : https://github.com/nod-ai/SHARK/blob/91ab594744ffbe982ababe340fcd208923ecae48/apps/stable_diffusion/web/ui/stablelm_ui.py#L61

With :

vicuna_model = UnshardedVicuna(

[Abhishek-Varma]

Replace this with safe_device = self.device.split("-")[0]

[Abhishek-Varma]

Replace this with :-

def compile(self):
        # Cannot load both the models in the memory at once
        # due to memory constraints, hence on demand compilation
        # is being used until the space is enough for both models

        # Testing : DO NOT Download Vmfbs if not found. Modify later
        # download vmfbs for A100
        supported_devices = ["cuda", "cpu-sync", "cpu-task", "cpu"]
        if (
            not self.first_vicuna_vmfb_path.exists()
            and self.device in supported_devices
            and self.precision in ["fp32", "fp16", "int8"]
        ):
            if (self.device == "cuda" and self.precision == "fp16") or (
                self.device in ["cpu-sync", "cpu-task"]
                and self.precision == "int8"
            ):
                download_public_file(
                    f"gs://shark_tank/vicuna/unsharded/vmfb/{self.first_vicuna_vmfb_path.name}",
                    self.first_vicuna_vmfb_path.absolute(),
                    single_file=True,
                )
            else:
                pass

        else:
            # get first vic
            # TODO: Remove after testing to avoid memory overload
            # fvic_shark_model = self.compile_first_vicuna()
            pass
        if (
            not self.second_vicuna_vmfb_path.exists()
            and self.device in supported_devices
            and self.precision in ["fp32", "fp16", "int8"]
        ):
            if (self.device == "cuda" and self.precision == "fp16") or (
                self.device in ["cpu-sync", "cpu-task"]
                and self.precision == "int8"
            ):
                download_public_file(
                    f"gs://shark_tank/vicuna/unsharded/vmfb/{self.second_vicuna_vmfb_path.name}",
                    self.second_vicuna_vmfb_path.absolute(),
                    single_file=True,
                )
            else:
                pass
        else:
            # get second vic
            # TODO: Remove after testing to avoid memory overload
            # svic_shark_model = self.compile_second_vicuna()
            pass

        return None
        # return tuple of shark_modules once mem is supported
        # return fvic_shark_model, svic_shark_model