pretrain fixes for mm

docs/llava.md

@@ -0,0 +1,36 @@
+# LLaVA
+
+### Installing dependencies
+
+```shell
+git clone https://github.com/haotian-liu/LLaVA.git
+cd LLaVA
+pip install --no-deps -e .
+```
+
+### Downloading assets
+
+LLaVA doesn't support remote datasets, so both the JSON and image assets need to be downloaded locally
+
+```shell
+mkdir llava
+mkdir data
+cd llava
+curl -L -O https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/resolve/main/images.zip
+unzip images.zip
+
+cd ../data
+curl -L -O https://huggingface.co/datasets/liuhaotian/LLaVA-Pretrain/resolve/main/blip_laion_cc_sbu_558k.json
+```
+
+### Pretraining
+
+Pretraining aligns the vision model with the language model.
+
+```shell
+accelerate launch -m axolotl.cli.train_mm examples/multimodal/pretrain-llava-llama.yml
+```
+
+### Finetuning
+
+TBD

examples/multimodal/pretrain-llava-llama.yml

@@ -0,0 +1,66 @@
+base_model: NousResearch/Llama-2-7b-hf
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+is_llama_derived_model: true
+
+# multimodal pretrain
+multimodal: true
+mm_vision_tower: openai/clip-vit-large-patch14
+tune_mm_mlp_adapter: true
+mm_freeze_backbone: true
+mm_vision_select_layer: -2
+mm_projector_type: mlp2x_gelu
+mm_image_folder: ./llava/
+mm_use_im_patch_token: false
+
+load_in_8bit: false
+load_in_4bit: false
+strict: false
+
+datasets:
+  - path: ./data/blip_laion_cc_sbu_558k.json
+dataset_prepared_path:
+val_set_size: 0.0
+output_dir: ./out
+
+sequence_len: 2048
+sample_packing: false
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_run_id:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 1
+optimizer: adamw_torch
+lr_scheduler: cosine
+learning_rate: 0.002
+
+train_on_inputs: false
+group_by_length: false
+bf16: true
+fp16: false
+tf32: false
+
+gradient_checkpointing: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+
+warmup_steps: 10
+eval_steps:
+save_steps: 0.1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:
+  pad_token: "<unk>"

examples/multimodal/pretrain-llava-mistral.yml

@@ -2,26 +2,29 @@ base_model: mistralai/Mistral-7B-v0.1
 model_type: MistralForCausalLM
 tokenizer_type: LlamaTokenizer
 is_mistral_derived_model: true
-multimodal: true
 
-vision_tower: openai/clip-vit-large-patch14
+# multimodal pretrain
+multimodal: true
+mm_vision_tower: openai/clip-vit-large-patch14
 tune_mm_mlp_adapter: true
+mm_freeze_backbone: true
 mm_vision_select_layer: -2
 mm_projector_type: mlp2x_gelu
 mm_image_folder: ./llava/
+mm_use_im_patch_token: false
 
 load_in_8bit: false
 load_in_4bit: false
 strict: false
 
 datasets:
-  - path: liuhaotian/LLaVA-CC3M-Pretrain-595K
+  - path: ./data/blip_laion_cc_sbu_558k.json
 dataset_prepared_path:
-val_set_size: 0.01
+val_set_size: 0.0
 output_dir: ./out
 
 sequence_len: 2048
-sample_packing: true
+sample_packing: false
 pad_to_sequence_len: true
 
 wandb_project:
@@ -32,8 +35,8 @@ wandb_log_model:
 
 gradient_accumulation_steps: 4
 micro_batch_size: 2
-num_epochs: 4
-optimizer: adamw_bnb_8bit
+num_epochs: 1
+optimizer: adamw_torch
 lr_scheduler: cosine
 learning_rate: 0.002
 
@@ -52,7 +55,7 @@ xformers_attention:
 flash_attention: true
 
 warmup_steps: 10
-eval_steps: 0.05
+eval_steps:
 save_steps:
 debug:
 deepspeed:

src/axolotl/cli/__init__.py

@@ -237,10 +237,11 @@ def load_mm_dataset(
         image_grid_pinpoints=cfg.mm_image_grid_pinpoints or None,
     )
     data_args.image_processor = vision_tower.image_processor
+    data_args.mm_use_im_start_end = cfg.mm_use_im_start_end or False
     tokenizer = load_tokenizer(cfg)
     train_dataset = LazySupervisedDataset(
         tokenizer=tokenizer,
-        data_path=data_args["data_path"],
+        data_path=data_args.data_path,
         data_args=data_args,
     )
 

src/axolotl/cli/train_mm.py

@@ -5,6 +5,7 @@ import logging
 from pathlib import Path
 
 import fire
+import torch
 import transformers
 from colorama import Fore
 
@@ -47,6 +48,8 @@ def do_cli(config: Path = Path("examples/"), **kwargs):
     dataset_meta = load_mm_dataset(
         cfg=parsed_cfg, cli_args=parsed_cli_args, model=model
     )
+    del model
+    torch.cuda.empty_cache()
     if parsed_cli_args.prepare_ds_only:
         return
     train(cfg=parsed_cfg, cli_args=parsed_cli_args, dataset_meta=dataset_meta)

src/axolotl/core/trainer_builder.py

@@ -110,6 +110,17 @@ class AxolotlTrainingArguments(TrainingArguments):
     bench_source_max_len: int = field(
         default=2048, metadata={"help": "Maximum source sequence length for bench."}
     )
+    tune_mm_mlp_adapter: bool = field(
+        default=False,
+        metadata={"help": "Whether to train the multimodal projector adapter"},
+    )
+    freeze_mm_mlp_adapter: bool = field(
+        default=False,
+        metadata={"help": "Whether to freeze the multimodal projector adapter"},
+    )
+    mm_projector_lr: Optional[float] = field(
+        default=None, metadata={"help": "Learning rate for the multimodal projector"}
+    )
 
 
 class AxolotlTrainer(Trainer):
@@ -260,21 +271,26 @@ class AxolotlTrainer(Trainer):
             run_dir = self._get_output_dir(trial=trial)
             output_dir = os.path.join(run_dir, checkpoint_folder)
 
-            # Only save Adapter
-            keys_to_match = ["mm_projector", "vision_resampler"]
-            if getattr(self.args, "use_im_start_end", False):
-                keys_to_match.extend(["embed_tokens", "embed_in"])
-
-            weight_to_save = get_mm_adapter_state_maybe_zero_3(
-                self.model.named_parameters(), keys_to_match
-            )
+            weights_to_save = self._get_mm_mlp_adapter_weights()
 
             if self.args.local_rank in (0, -1):
                 self.model.config.save_pretrained(output_dir)
-                torch.save(weight_to_save, os.path.join(output_dir, "mm_projector.bin"))
+                torch.save(
+                    weights_to_save, os.path.join(output_dir, "mm_projector.bin")
+                )
         else:
             super()._save_checkpoint(model, trial, metrics)
 
+    def _get_mm_mlp_adapter_weights(self):
+        # Only save Adapter
+        keys_to_match = ["mm_projector", "vision_resampler"]
+        if getattr(self.args, "use_im_start_end", False):
+            keys_to_match.extend(["embed_tokens", "embed_in"])
+
+        return get_mm_adapter_state_maybe_zero_3(
+            self.model.named_parameters(), keys_to_match
+        )
+
     def _save(self, output_dir: Optional[str] = None, state_dict=None):
         if getattr(self.args, "tune_mm_mlp_adapter", False):
             pass
@@ -648,8 +664,17 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         training_arguments_kwargs[
             "sample_packing_seq_len_multiplier"
         ] = self.cfg.micro_batch_size
+
         training_arguments_kwargs["relora_steps"] = self.cfg.relora_steps
         training_arguments_kwargs["relora_warmup_steps"] = self.cfg.relora_warmup_steps
+
+        # multimodal: llava
+        training_arguments_kwargs["tune_mm_mlp_adapter"] = self.cfg.tune_mm_mlp_adapter
+        training_arguments_kwargs[
+            "freeze_mm_mlp_adapter"
+        ] = self.cfg.freeze_mm_mlp_adapter
+        training_arguments_kwargs["mm_projector_lr"] = self.cfg.mm_projector_lr
+
         training_arguments_kwargs = self.hook_pre_create_training_args(
             training_arguments_kwargs
         )

src/axolotl/train.py

@@ -159,14 +159,14 @@ def train(
         # The model name saved is `pytorch_model.bin`
         unwrapped_model.save_pretrained(
             cfg.output_dir,
-            is_main_process=trainer.accelerator.is_main_process,
+            is_main_process=trainer.args.should_save,
             save_function=trainer.accelerator.save,
             state_dict=trainer.accelerator.get_state_dict(trainer.model_wrapped),
         )
-    elif cfg.local_rank == 0:
+    elif trainer.args.should_save:
         if cfg.flash_optimum:
             model = BetterTransformer.reverse(model)
-
+        # TODO figure out if `trainer.save_model(cfg.output_dir)` is sufficient here
         model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
 
     if not cfg.hub_model_id:

src/axolotl/utils/models.py

@@ -278,12 +278,21 @@ def load_model(
             if cfg.mm_freeze_backbone:
                 model.model.requires_grad_(False)
 
-            def make_inputs_require_grad(
-                module, input, output
-            ):  # pylint: disable=redefined-builtin,unused-argument
-                output.requires_grad_(True)
+            if cfg.gradient_checkpointing:
+                if hasattr(model, "enable_input_require_grads"):
+                    model.enable_input_require_grads()
+                else:
 
-            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+                    def make_inputs_require_grad(
+                        module,
+                        input,
+                        output,
+                    ):  # pylint: disable=redefined-builtin,unused-argument
+                        output.requires_grad_(True)
+
+                    model.get_input_embeddings().register_forward_hook(
+                        make_inputs_require_grad
+                    )
 
             model_args = ModelArguments(
                 model_name_or_path=cfg.base_model,
@@ -295,17 +304,17 @@ def load_model(
                 pretrain_mm_mlp_adapter=cfg.pretrain_mm_mlp_adapter,
                 mm_projector_type=cfg.mm_projector_type or "linear",
                 mm_use_im_start_end=cfg.mm_use_im_start_end or False,
-                mm_use_im_patch_token=cfg.mm_use_im_patch_token or True,
+                mm_use_im_patch_token=cfg.mm_use_im_patch_token,
                 mm_vision_select_feature=cfg.mm_vision_select_feature or "patch",
             )
 
-            if cfg.mm_vision_tower:
+            if cfg.mm_vision_tower is not None:
                 model.get_model().initialize_vision_modules(
                     model_args=model_args, fsdp=cfg.fsdp
                 )
 
             vision_tower = model.get_vision_tower()
-            vision_tower.to(dtype=cfg.torch_dtype)
+            vision_tower.to(dtype=cfg.torch_dtype, device=cfg.device)
 
             # pylint: disable=duplicate-code
             data_args = DataArguments(
@@ -321,8 +330,8 @@ def load_model(
             data_args.image_processor = vision_tower.image_processor
             model.config.image_aspect_ratio = data_args.image_aspect_ratio
             model.config.image_grid_pinpoints = data_args.image_grid_pinpoints
-            model.config.tune_mm_mlp_adapter = model_args.tune_mm_mlp_adapter
-            if model_args.tune_mm_mlp_adapter:
+            model.config.tune_mm_mlp_adapter = cfg.tune_mm_mlp_adapter
+            if cfg.tune_mm_mlp_adapter:
                 model.requires_grad_(False)
                 for (
                     p  # pylint: disable=invalid-name
@@ -338,8 +347,8 @@ def load_model(
 
             model.config.mm_use_im_start_end = (
                 data_args.mm_use_im_start_end
-            ) = model_args.mm_use_im_start_end
-            model.config.mm_use_im_patch_token = model_args.mm_use_im_patch_token
+            ) = cfg.mm_use_im_start_end
+            model.config.mm_use_im_patch_token = cfg.mm_use_im_patch_token
             model.initialize_vision_tokenizer(model_args, tokenizer=tokenizer)
         elif cfg.is_llama_derived_model and not cfg.trust_remote_code and not cfg.gptq:
             from transformers import LlamaForCausalLM

src/axolotl/utils/trainer.py

@@ -13,7 +13,7 @@ import torch.distributed as dist
 from datasets import set_caching_enabled
 from torch.utils.data import DistributedSampler, RandomSampler
 
-from axolotl.core.trainer_builder import HFCausalTrainerBuilder
+from axolotl.core.trainer_builder import AxolotlTrainer, HFCausalTrainerBuilder
 from axolotl.utils.collators import DataCollatorForSeq2Seq
 from axolotl.utils.dataloader import MultipackDistributedDataloader
 from axolotl.utils.distributed import (
@@ -259,7 +259,9 @@ def setup_fsdp_envs(cfg):
         ] = cfg.fsdp_config.fsdp_transformer_layer_cls_to_wrap
 
 
-def setup_trainer(cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps):
+def setup_trainer(
+    cfg, train_dataset, eval_dataset, model, tokenizer, total_num_steps
+) -> AxolotlTrainer:
     if cfg.fsdp:
         setup_fsdp_envs(cfg)
     elif cfg.deepspeed: