quant_configs_and_filter_fns = [
    (QuantConfig(bitwidth=2, group_size=128), use_2bit),
    (QuantConfig(bitwidth=3, group_size=256), use_3bit),
    (QuantConfig(bitwidth=4, group_size=256), use_4bit),
    (QuantConfig(bitwidth=8), use_8bit),
 ]

optimizer = create_optimizer(model, quant_configs_and_filter_fns, base_optimizer_cls, base_optimizer_kwargs)

def use_2bit(m, fqn):
        if not fqn.endswith(".weight"):
            return False

        # Layers 0-3
        for i in range(4):
            if fqn.startswith(f"model.layers.{i}."):
                for key in ["v_proj", "gate_proj", "up_proj"]:
                    if key in fqn:
                        return True

        # Layers 4-31
        for i in range(4, 32):
            if fqn.startswith(f"model.layers.{i}."):
                for key in ["q_proj", "v_proj", "gate_proj", "up_proj", "down_proj"]:
                    if key in fqn:
                        return True

        return False