nexaai 1.0.29__cp310-cp310-macosx_14_0_universal2.whl

nexaai/mlx_backend/vlm/modeling/trainer/trainer.py

@@ -0,0 +1,296 @@
+import json
+import warnings
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Union
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mlx.utils import tree_flatten, tree_map
+
+
+def get_prompt(model_type, processor, conversation):
+    if model_type == "paligemma":
+        return conversation
+
+    if "chat_template" in processor.__dict__.keys():
+        prompt = processor.apply_chat_template(
+            conversation,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+    elif "tokenizer" in processor.__dict__.keys():
+        prompt = processor.tokenizer.apply_chat_template(
+            conversation,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+    return prompt
+
+
+class Dataset:
+    def __init__(
+        self,
+        hf_dataset,
+        config,
+        processor,
+        image_processor=None,
+        take=None,
+        split=None,
+        image_resize_shape=None,
+    ):
+        if split is not None:
+            self.dataset = hf_dataset[split]
+        else:
+            self.dataset = hf_dataset
+        if take is not None:
+            self.dataset = self.dataset.take(take)
+        self.processor = processor
+        self.config = config
+        self.image_processor = image_processor
+        self.image_resize_shape = image_resize_shape
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        from mlx_vlm.utils import prepare_inputs
+
+        item = self.dataset[idx]
+
+        images = item["images"]
+        conversations = item["messages"]
+        prompts = []
+
+        if isinstance(conversations, list) and isinstance(conversations[0], list):
+            for conversation in conversations:
+                if self.config["model_type"] == "pixtral":
+                    conversation = [json.loads(i) for i in conversation]
+                    if len(conversations) > 1:
+                        warnings.warn(
+                            "Pixtral batch processing is not supported yet. Set batch size to 1."
+                        )
+
+                prompt = get_prompt(
+                    self.config["model_type"], self.processor, conversation
+                )
+                prompts.append(prompt)
+
+        else:
+            if self.config["model_type"] == "pixtral":
+                conversations = [json.loads(i) for i in conversations]
+            prompt = get_prompt(
+                self.config["model_type"], self.processor, conversations
+            )
+            prompts.append(prompt)
+
+        image_token_index = self.config["image_token_index"]
+
+        inputs = prepare_inputs(
+            self.processor,
+            images,
+            prompts,
+            image_token_index,
+            self.image_resize_shape,
+        )
+        input_ids = inputs["input_ids"]
+        pixel_values = inputs["pixel_values"]
+        mask = inputs["attention_mask"]
+        kwargs = {
+            k: v
+            for k, v in inputs.items()
+            if k not in ["input_ids", "pixel_values", "attention_mask"]
+        }
+
+        if mask is None:
+            mask = mx.ones_like(input_ids)
+
+        return {
+            "pixel_values": pixel_values,
+            "input_ids": input_ids,
+            "attention_mask": mask,
+            **kwargs,
+        }
+
+
+def grad_checkpoint(layer):
+    """
+    Update all instances of type(layer) to use gradient checkpointing.
+    """
+    fn = type(layer).__call__
+
+    def checkpointed_fn(model, *args, **kwargs):
+        def inner_fn(params, *args, **kwargs):
+            model.update(params)
+            return fn(model, *args, **kwargs)
+
+        return mx.checkpoint(inner_fn)(model.trainable_parameters(), *args, **kwargs)
+
+    type(layer).__call__ = checkpointed_fn
+
+
+@dataclass
+class TrainingArgs:
+    batch_size: int = field(default=4, metadata={"help": "Minibatch size."})
+    iters: int = field(default=100, metadata={"help": "Iterations to train for."})
+    val_batches: int = field(
+        default=25,
+        metadata={
+            "help": "Number of validation batches, -1 uses the entire validation set."
+        },
+    )
+    steps_per_report: int = field(
+        default=10,
+        metadata={"help": "Number of training steps between loss reporting."},
+    )
+    steps_per_eval: int = field(
+        default=200, metadata={"help": "Number of training steps between validations."}
+    )
+    steps_per_save: int = field(
+        default=100, metadata={"help": "Save the model every number steps"}
+    )
+    max_seq_length: int = field(
+        default=2048, metadata={"help": "Maximum sequence length."}
+    )
+    adapter_file: str = field(
+        default="adapters.safetensors",
+        metadata={"help": "Save/load path for the trained adapter weights."},
+    )
+    grad_checkpoint: bool = field(
+        default=False,
+        metadata={"help": "Use gradient checkpointing to reduce memory use."},
+    )
+
+
+def default_loss(model, inputs, targets, lengths):
+    logits = model(inputs)
+    logits = logits.astype(mx.float32)
+
+    length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]
+
+    ce = nn.losses.cross_entropy(logits, targets) * length_mask
+    ntoks = length_mask.sum()
+    ce = ce.sum() / ntoks
+
+    return ce, ntoks
+
+
+class Trainer:
+    def __init__(
+        self,
+        model,
+        optimizer,
+        train_on_completions=False,
+        assistant_id=77091,
+        clip_gradients=None,
+    ):
+        self.model = model
+        self.optimizer = optimizer
+        self.train_on_completions = train_on_completions
+        self.assistant_id = assistant_id
+        self.clip_gradients = clip_gradients
+
+    def loss_fn(self, model, batch):
+        pixel_values = batch["pixel_values"]
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        lengths = mx.sum(attention_mask, axis=1)
+        labels = input_ids[:, 1:]
+
+        batch_size, seq_length = input_ids.shape
+
+        if self.train_on_completions:
+            weight_mask = mx.ones_like(attention_mask)
+
+            assistant_response_index = np.where(input_ids == self.assistant_id)[1]
+            range_matrix = mx.repeat(
+                mx.expand_dims(mx.arange(seq_length), 0), batch_size, axis=0
+            )
+            assistant_mask = range_matrix <= mx.array(assistant_response_index).reshape(
+                -1, 1
+            )
+            # Apply the mask to weight_mask
+            weight_mask = mx.where(
+                assistant_mask, mx.zeros_like(weight_mask), weight_mask
+            )[:, 1:]
+        else:
+            weight_mask = None
+
+        input_ids = input_ids[:, :-1]
+
+        kwargs = {
+            k: v
+            for k, v in batch.items()
+            if k not in ["input_ids", "pixel_values", "attention_mask"]
+        }
+
+        # Forward pass
+        outputs = model(input_ids, pixel_values, attention_mask, **kwargs)
+
+        # Cast to float32
+        logits = outputs.logits.astype(mx.float32)
+
+        # Ensure logits and labels have the same sequence length
+        def align_logits_with_labels(logits, labels):
+            if logits.shape[1] < labels.shape[1]:
+                pad_length = labels.shape[1] - logits.shape[1]
+                pad_width = ((0, 0), (0, pad_length), (0, 0))
+                return mx.pad(logits, pad_width, mode="constant", constant_values=-100)
+            elif logits.shape[1] > labels.shape[1]:
+                return logits[:, -labels.shape[1] :, :]
+            return logits
+
+        logits = align_logits_with_labels(logits, labels)
+
+        length_mask = mx.arange(input_ids.shape[1])[None, :] < lengths[:, None]
+
+        # Compute loss only on non-padded tokens
+        ce = (
+            nn.losses.cross_entropy(
+                logits,
+                labels,
+                weights=weight_mask,
+            )
+            * length_mask
+        )
+        ntoks = length_mask.sum()
+        ce = ce.sum() / ntoks
+
+        return ce
+
+    def train_step(self, batch):
+        loss_and_grad_fn = nn.value_and_grad(self.model, self.loss_fn)
+        loss, grads = loss_and_grad_fn(self.model, batch)
+
+        # Add gradient clipping
+        if self.clip_gradients is not None:
+            grads = tree_map(
+                lambda g: mx.clip(g, -self.clip_gradients, self.clip_gradients), grads
+            )
+
+        self.optimizer.update(self.model, grads)
+
+        return loss
+
+    @mx.compile
+    def train_epoch(self, dataloader):
+        total_loss = 0
+        for batch in dataloader:
+            loss = self.train_step(batch)
+            mx.eval(self.model, self.optimizer.state)
+            total_loss += loss
+        return total_loss / len(dataloader)
+
+
+def save_adapter(
+    model: nn.Module,
+    adapter_file: Union[str, Path],
+):
+    path = Path(adapter_file)
+    if hasattr(model.config, "lora"):
+        with open(path.parent / "adapter_config.json", "w") as f:
+            json.dump(model.config.lora, f)
+    flattened_tree = tree_flatten(model.trainable_parameters())
+    mx.save_safetensors(str(adapter_file), dict(flattened_tree))

nexaai/mlx_backend/vlm/modeling/trainer/utils.py

@@ -0,0 +1,160 @@
+import json
+from pathlib import Path
+
+import mlx.nn as nn
+from mlx.utils import tree_flatten
+
+from .lora import LoRaLayer
+
+
+def get_module_by_name(model, name):
+    parts = name.split(".")
+    module = model
+    for part in parts:
+        if part.isdigit():
+            module = module[int(part)]
+        else:
+            module = getattr(module, part)
+    return module
+
+
+def set_module_by_name(model, name, new_module):
+    parts = name.split(".")
+    module = model
+    for part in parts[:-1]:
+        if part.isdigit():
+            module = module[int(part)]
+        else:
+            module = getattr(module, part)
+    if parts[-1].isdigit():
+        module[int(parts[-1])] = new_module
+    else:
+        setattr(module, parts[-1], new_module)
+
+
+def get_peft_model(
+    model, linear_layers, rank=10, alpha=0.1, dropout=0.1, freeze=True, verbose=True
+):
+    if freeze:
+        freeze_model(model)
+
+    for name, module in model.language_model.named_modules():
+        if isinstance(module, nn.Linear) or isinstance(module, nn.QuantizedLinear):
+            if name.split(".")[-1] in linear_layers:
+                lora_layer = LoRaLayer(module, rank, alpha, dropout)
+                set_module_by_name(model.language_model, name, lora_layer)
+
+    model.config.lora = {}
+    model.config.lora["rank"] = rank
+    model.config.lora["alpha"] = alpha
+    model.config.lora["dropout"] = dropout
+
+    if verbose:
+        print_trainable_parameters(model.language_model)
+
+    return model
+
+
+def freeze_model(model):
+    for name, module in model.named_modules():
+        name = name.split(".")[0]
+        if name in [
+            "language_model",
+            "vision_model",
+            "vision_tower",
+            "aligner",
+            "connector",
+            "multi_modal_projector",
+            "mm_projector",
+        ]:
+            model[f"{name}"].freeze()
+
+
+def find_all_linear_names(model):
+    cls = nn.Linear
+    quantized_cls = nn.QuantizedLinear
+    lora_module_names = set()
+    multimodal_keywords = [
+        "mm_projector",
+        "vision_tower",
+        "vision_resampler",
+        "aligner",
+    ]
+    for name, module in model.named_modules():
+        if any(mm_keyword in name for mm_keyword in multimodal_keywords):
+            continue
+        if isinstance(module, cls) or isinstance(module, quantized_cls):
+            names = name.split(".")
+            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
+
+    if "lm_head" in lora_module_names:  # needed for 16-bit
+        lora_module_names.remove("lm_head")
+    return list(lora_module_names)
+
+
+def count_parameters(model):
+    def nparams(m):
+        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
+            return m.weight.size * (32 // m.bits)
+        return sum(v.size for _, v in tree_flatten(m.parameters()))
+
+    leaf_modules = tree_flatten(
+        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
+    )
+    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6
+
+    return total_p
+
+
+def print_trainable_parameters(model):
+    def nparams(m):
+        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
+            return m.weight.size * (32 // m.bits)
+        return sum(v.size for _, v in tree_flatten(m.parameters()))
+
+    leaf_modules = tree_flatten(
+        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
+    )
+    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6
+    trainable_p = (
+        sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 10**6
+    )
+
+    print(
+        f"#trainable params: {trainable_p} M || all params: {total_p} M || trainable%: {(trainable_p * 100 / total_p):.3f}%"
+    )
+
+
+def apply_lora_layers(model: nn.Module, adapter_path: str) -> nn.Module:
+    """
+    Apply LoRA layers to the model.
+
+    Args:
+        model (nn.Module): The neural network model.
+        adapter_path (str): Path to the adapter configuration file.
+
+    Returns:
+        nn.Module: The updated model with LoRA layers applied.
+    """
+    adapter_path = Path(adapter_path)
+
+    if not adapter_path.exists():
+        raise FileNotFoundError(f"The adapter path does not exist: {adapter_path}")
+
+    # Check if the adapter has lora params in the config (adapter_config.json)
+    with open(adapter_path / "adapter_config.json", "r") as f:
+        config = json.load(f)
+        if "rank" not in config:
+            raise ValueError("The adapter does not have lora params in the config")
+
+    # TODO: add lora params to the config and load them here
+    list_of_modules = find_all_linear_names(model.language_model.model)
+    if config is not None:
+        model = get_peft_model(model, list_of_modules, **config)
+    else:
+        model = get_peft_model(model, list_of_modules)
+
+    # TODO: Use custom adapter name
+    model.load_weights(str(adapter_path / "adapters.safetensors"), strict=False)
+
+    return model