Merge branch 'ml-explore:main' into adding-GRPO-training

llms/mlx_lm/convert.py

@@ -1,8 +1,27 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+from enum import Enum
 
-from .utils import convert
+from .utils import convert, mixed_2_6, mixed_3_6
+
+
+class MixedQuants(Enum):
+    mixed_3_6 = "mixed_3_6"
+    mixed_2_6 = "mixed_2_6"
+
+    @classmethod
+    def recipe_names(cls):
+        return [member.name for member in cls]
+
+
+def quant_args(arg):
+    try:
+        return MixedQuants[arg].value
+    except KeyError:
+        raise argparse.ArgumentTypeError(
+            f"Invalid q-recipe {arg!r}. Choose from: {MixedQuants.recipe_names()}"
+        )
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -29,6 +48,12 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
+    parser.add_argument(
+        "--quant-predicate",
+        help=f"Mixed-bit quantization recipe. Choices: {MixedQuants.recipe_names()}",
+        type=quant_args,
+        required=False,
+    )
     parser.add_argument(
         "--dtype",
         help="Type to save the non-quantized parameters.",

llms/mlx_lm/evaluate.py

@@ -289,17 +289,15 @@ class MLXLM(LM):
         contexts, options = zip(*[req.args for req in requests])
         # contrary to the doc the second element of the tuple contains
         # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
-        keys = list(options[0].keys())
-        assert "until" in keys
-        untils = [x["until"] for x in options]
         completions = []
 
-        for context, until in tqdm(zip(contexts, untils), total=len(contexts)):
+        for context, opt in tqdm(zip(contexts, options), total=len(contexts)):
+            until = opt["until"]
             context = self.tokenizer.encode(
                 context, add_special_tokens=not self.use_chat_template
             )
             max_tokens = min(
-                self._max_tokens,
+                opt.get("max_gen_tokens", self._max_tokens),
                 self.tokenizer.model_max_length - len(context),
             )
             text = ""
@@ -334,9 +332,9 @@ def main():
     )
     parser.add_argument(
         "--limit",
-        default=1.0,
+        default=100,
         help="Limit the number of examples per task.",
-        type=float,
+        type=int,
     )
     parser.add_argument("--seed", type=int, default=123, help="Random seed.")
     parser.add_argument(

llms/mlx_lm/manage.py

@@ -2,7 +2,22 @@ import argparse
 from typing import List, Union
 
 from huggingface_hub import scan_cache_dir
-from transformers.commands.user import tabulate
+
+
+def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
+    """
+    Inspired by:
+    - stackoverflow.com/a/8356620/593036
+    - stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
+    """
+    col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
+    row_format = ("{{:{}}} " * len(headers)).format(*col_widths)
+    lines = []
+    lines.append(row_format.format(*headers))
+    lines.append(row_format.format(*["-" * w for w in col_widths]))
+    for row in rows:
+        lines.append(row_format.format(*row))
+    return "\n".join(lines)
 
 
 def ask_for_confirmation(message: str) -> bool:

llms/mlx_lm/models/phi3.py

@@ -23,8 +23,10 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float = 10000
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, List[float]]]] = None
+    partial_rotary_factor: float = 1.0
     max_position_embeddings: int = 131072
     original_max_position_embeddings: int = 4096
+    tie_word_embeddings: bool = False
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
@@ -59,9 +61,10 @@ class Attention(nn.Module):
         self.qkv_proj = nn.Linear(dim, op_size, bias=False)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
 
+        rope_dim = int(head_dim * args.partial_rotary_factor)
         if args.rope_scaling and args.rope_scaling["type"] in ["longrope", "su"]:
             self.rope = SuScaledRotaryEmbedding(
-                head_dim,
+                rope_dim,
                 base=args.rope_theta,
                 max_position_embeddings=args.max_position_embeddings,
                 original_max_position_embeddings=args.original_max_position_embeddings,
@@ -74,7 +77,7 @@ class Attention(nn.Module):
                 assert isinstance(args.rope_scaling["factor"], float)
                 rope_scale = 1 / args.rope_scaling["factor"]
             self.rope = nn.RoPE(
-                head_dim,
+                rope_dim,
                 traditional=args.rope_traditional,
                 base=args.rope_theta,
                 scale=rope_scale,
@@ -190,7 +193,8 @@ class Model(nn.Module):
         super().__init__()
         self.model_type = args.model_type
         self.model = Phi3Model(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
         self.args = args
 
     def __call__(
@@ -200,7 +204,11 @@ class Model(nn.Module):
         cache=None,
     ):
         out = self.model(inputs, mask, cache)
-        return self.lm_head(out)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
 
     @property
     def layers(self):

llms/mlx_lm/models/su_rope.py

@@ -51,11 +51,13 @@ class SuScaledRotaryEmbedding(nn.Module):
             + math.log(max_position_embeddings / original_max_position_embeddings)
             / math.log(original_max_position_embeddings)
         )
+        self.dim = dims
 
     def __call__(self, x, offset: int = 0):
+        x[..., : self.dim] = self.scale * x[..., : self.dim]
         return mx.fast.rope(
-            self.scale * x,
-            x.shape[-1],
+            x,
+            self.dim,
             traditional=False,
             base=None,
             scale=1.0,

llms/mlx_lm/utils.py

@@ -1015,6 +1015,46 @@ def save_config(
         json.dump(config, fid, indent=4)
 
 
+def mixed_quant_predicate_builder(
+    low_bits: int = 4, high_bits: int = 4, group_size: int = 64
+) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+    def mixed_quant_predicate(
+        path: str,
+        module: nn.Module,
+        config: dict,
+    ) -> Union[bool, dict]:
+        """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
+        Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
+        By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
+        """
+
+        if not hasattr(module, "to_quantized"):
+            return False
+
+        index = int(path.split(".")[2]) if len(path.split(".")) > 2 else 0
+
+        num_layers = config["num_hidden_layers"]
+        use_more_bits = (
+            index < num_layers // 8
+            or index >= 7 * num_layers // 8
+            or (index - num_layers // 8) % 3 == 2
+        )
+        if "v_proj" in path and use_more_bits:
+            return {"group_size": group_size, "bits": high_bits}
+        if "down_proj" in path and use_more_bits:
+            return {"group_size": group_size, "bits": high_bits}
+        if "lm_head" in path:
+            return {"group_size": group_size, "bits": high_bits}
+
+        return {"group_size": group_size, "bits": low_bits}
+
+    return mixed_quant_predicate
+
+
+mixed_3_6 = mixed_quant_predicate_builder(low_bits=3)
+mixed_2_6 = mixed_quant_predicate_builder(low_bits=2)
+
+
 def convert(
     hf_path: str,
     mlx_path: str = "mlx_model",