From 15d013aa5f9acebdba4dd50806d185f1b1935f89 Mon Sep 17 00:00:00 2001
From: Elisei Rykov <e.rykov@tbank.ru>
Date: Sat, 28 Sep 2024 15:44:15 +0300
Subject: [PATCH] in progress

---
 .../multimodal/attention_based_selector.json  | 205 ++++++++++++++++++
 .../train/multimodal/attention_pooling.json   | 205 ++++++++++++++++++
 configs/exp/train/multimodal/llava.json       | 205 ++++++++++++++++++
 .../train/multimodal/llava_non_pickled.json   | 205 ++++++++++++++++++
 .../accelerate_llava_next_data.json           |  14 ++
 configs/utils/preprocess/coco2014_clip.json   |  13 +-
 configs/utils/preprocess/llava_next_data.json |  14 ++
 .../common/data/multimodal/common.py          |   3 +-
 .../dataset/multimodal/collators.py           |   2 +-
 9 files changed, 859 insertions(+), 7 deletions(-)
 create mode 100644 configs/exp/train/multimodal/attention_based_selector.json
 create mode 100644 configs/exp/train/multimodal/attention_pooling.json
 create mode 100644 configs/exp/train/multimodal/llava.json
 create mode 100644 configs/exp/train/multimodal/llava_non_pickled.json
 create mode 100644 configs/utils/preprocess/accelerate_llava_next_data.json
 create mode 100644 configs/utils/preprocess/llava_next_data.json

diff --git a/configs/exp/train/multimodal/attention_based_selector.json b/configs/exp/train/multimodal/attention_based_selector.json
new file mode 100644
index 0000000..b22f3c3
--- /dev/null
+++ b/configs/exp/train/multimodal/attention_based_selector.json
@@ -0,0 +1,205 @@
+{
+    "train_dataset_settings": {
+        "sources": [
+            {
+                "name": "train",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/train_chat.jsonl",
+                "sample_rate": 1.0
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+        "n_modality_embeddings": 576,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "val_dataset_settings": {
+        "sources": [
+            {
+                "name": "test",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                "num_samples": 5000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+            },
+        "n_modality_embeddings": 576,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "model_settings": {
+        "model_path": "/from_s3/model",
+        "model_type": "causal",
+        "transformers_settings": {
+        },
+        "model_kwargs": {
+            "attn_implementation": "flash_attention_2"
+        },
+        "embeddings_initialization_strategy": {
+            "<MS>": "bot",
+            "</MS>": "bot",
+            "<img>": "bot",
+            "<audio>": "bot"
+        },
+        "peft_settings": {
+            "r": 16,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "target_modules": [
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "o_proj"
+            ],
+            "task_type": "CAUSAL_LM",
+            "modules_to_save": ["embed_tokens", "lm_head"],
+            "name": "LORA"
+        }
+    },
+    "tokenizer_settings": {
+        "tokenizer_path": "/from_s3/model"
+    },
+    "special_tokens_settings": {
+        "bos_token": "<|begin_of_text|>",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|end_of_text|>"
+    },
+    "trainer_settings": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "eval_steps": 300,
+        "save_steps": 300,
+        "per_device_train_batch_size": 2,
+        "per_device_eval_batch_size": 2,
+        "gradient_accumulation_steps": 16,
+        "logging_steps": 1,
+        "learning_rate": 1e-4,
+        "num_train_epochs": 1,
+        "lr_scheduler_type": "cosine",
+        "warmup_ratio": 0.3,
+        "fp16": false,
+        "bf16": true,
+        "optim": "adamw_torch",
+        "save_total_limit": 5,
+        "load_best_model_at_end": false,
+        "deepspeed": "configs/exp/deepspeed/stage2.json",
+        "dispatch_batches": false,
+        "dataloader_num_workers": 16,
+        "dataloader_pin_memory": false
+    },
+    "wandb_settings": {
+        "project_name": "rykov",
+        "run_name": "multimodal",
+        "entity": "rykov"
+    },
+    "log_path": "train_output",
+    "modality_encoder_settings_mapping": {
+        "image": {
+            "modality_encoder_type": "clip",
+            "is_pickle": true,
+            "encoder_path": "/from_s3/clip"
+        },
+        "audio": null
+    },
+    "modality_projector_mapping": {
+        "image": "threshold_selector",
+        "audio": null
+    },
+    "modality_projector_initialization_mapping": {
+        "image": null,
+        "audio": null
+    },
+    "cherry_pick_settings": {
+        "generator_transformers_settings": {
+            "num_beams": 1,
+            "max_new_tokens": 64,
+            "repetition_penalty": 1.0,
+            "stop_strings": "<|eot_id|>"
+        },
+        "custom_generation_settings": {
+            "skip_special_tokens": true
+          },
+        "dataset_settings": {
+            "sources": [
+                {
+                    "name": "cherry_pickls",
+                    "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                    "num_samples": 50
+                }
+            ],
+            "prompt_template": {
+                "role_tag_mapping": {
+                    "bot": "assistant",
+                    "user": "user",
+                    "system": "system"
+                },
+                "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+                "suffix_template": "<|eot_id|>"
+            },
+            "dataset_type": "multimodal",
+            "max_tokens_count": 2000,
+            "n_modality_embeddings": 576,
+            "start_modality_token": "<MS>",
+            "end_modality_token": "</MS>",
+            "only_answer_loss": true,
+            "random_cut": true,
+            "modality_token_mapping": {
+                "image": "<img>",
+                "audio": "<audio>"
+            },
+            "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+            "truncate_top": false
+        },
+        "metric_settings": []
+    }
+}
diff --git a/configs/exp/train/multimodal/attention_pooling.json b/configs/exp/train/multimodal/attention_pooling.json
new file mode 100644
index 0000000..453561c
--- /dev/null
+++ b/configs/exp/train/multimodal/attention_pooling.json
@@ -0,0 +1,205 @@
+{
+    "train_dataset_settings": {
+        "sources": [
+            {
+                "name": "train",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/train_chat.jsonl",
+                "sample_rate": 1.0
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+        "n_modality_embeddings": 64,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "val_dataset_settings": {
+        "sources": [
+            {
+                "name": "test",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                "num_samples": 5000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+            },
+        "n_modality_embeddings": 64,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "model_settings": {
+        "model_path": "/from_s3/model",
+        "model_type": "causal",
+        "transformers_settings": {
+        },
+        "model_kwargs": {
+            "attn_implementation": "flash_attention_2"
+        },
+        "embeddings_initialization_strategy": {
+            "<MS>": "bot",
+            "</MS>": "bot",
+            "<img>": "bot",
+            "<audio>": "bot"
+        },
+        "peft_settings": {
+            "r": 16,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "target_modules": [
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "o_proj"
+            ],
+            "task_type": "CAUSAL_LM",
+            "modules_to_save": ["embed_tokens", "lm_head"],
+            "name": "LORA"
+        }
+    },
+    "tokenizer_settings": {
+        "tokenizer_path": "/from_s3/model"
+    },
+    "special_tokens_settings": {
+        "bos_token": "<|begin_of_text|>",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|end_of_text|>"
+    },
+    "trainer_settings": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "eval_steps": 400,
+        "save_steps": 400,
+        "per_device_train_batch_size": 2,
+        "per_device_eval_batch_size": 2,
+        "gradient_accumulation_steps": 16,
+        "logging_steps": 1,
+        "learning_rate": 1e-4,
+        "num_train_epochs": 1,
+        "lr_scheduler_type": "cosine",
+        "warmup_ratio": 0.3,
+        "fp16": false,
+        "bf16": true,
+        "optim": "adamw_torch",
+        "save_total_limit": 5,
+        "load_best_model_at_end": false,
+        "deepspeed": "configs/exp/deepspeed/stage2.json",
+        "dispatch_batches": false,
+        "dataloader_num_workers": 16,
+        "dataloader_pin_memory": false
+    },
+    "wandb_settings": {
+        "project_name": "rykov",
+        "run_name": "multimodal",
+        "entity": "rykov"
+    },
+    "log_path": "train_output",
+    "modality_encoder_settings_mapping": {
+        "image": {
+            "modality_encoder_type": "clip",
+            "is_pickle": true,
+            "encoder_path": "/from_s3/clip"
+        },
+        "audio": null
+    },
+    "modality_projector_mapping": {
+        "image": "top_k_attention_pooling",
+        "audio": null
+    },
+    "modality_projector_initialization_mapping": {
+        "image": null,
+        "audio": null
+    },
+    "cherry_pick_settings": {
+        "generator_transformers_settings": {
+            "num_beams": 1,
+            "max_new_tokens": 64,
+            "repetition_penalty": 1.0,
+            "stop_strings": "<|eot_id|>"
+        },
+        "custom_generation_settings": {
+            "skip_special_tokens": true
+          },
+        "dataset_settings": {
+            "sources": [
+                {
+                    "name": "cherry_pickls",
+                    "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                    "num_samples": 50
+                }
+            ],
+            "prompt_template": {
+                "role_tag_mapping": {
+                    "bot": "assistant",
+                    "user": "user",
+                    "system": "system"
+                },
+                "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+                "suffix_template": "<|eot_id|>"
+            },
+            "dataset_type": "multimodal",
+            "max_tokens_count": 2000,
+            "n_modality_embeddings": 64,
+            "start_modality_token": "<MS>",
+            "end_modality_token": "</MS>",
+            "only_answer_loss": true,
+            "random_cut": true,
+            "modality_token_mapping": {
+                "image": "<img>",
+                "audio": "<audio>"
+            },
+            "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+            "truncate_top": false
+        },
+        "metric_settings": []
+    }
+}
diff --git a/configs/exp/train/multimodal/llava.json b/configs/exp/train/multimodal/llava.json
new file mode 100644
index 0000000..2617be3
--- /dev/null
+++ b/configs/exp/train/multimodal/llava.json
@@ -0,0 +1,205 @@
+{
+    "train_dataset_settings": {
+        "sources": [
+            {
+                "name": "train",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/train_chat.jsonl",
+                "sample_rate": 1.0
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+        "n_modality_embeddings": 576,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "val_dataset_settings": {
+        "sources": [
+            {
+                "name": "test",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                "num_samples": 10000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+            },
+        "n_modality_embeddings": 576,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "model_settings": {
+        "model_path": "/from_s3/model",
+        "model_type": "causal",
+        "transformers_settings": {
+        },
+        "model_kwargs": {
+            "attn_implementation": "flash_attention_2"
+        },
+        "embeddings_initialization_strategy": {
+            "<MS>": "bot",
+            "</MS>": "bot",
+            "<img>": "bot",
+            "<audio>": "bot"
+        },
+        "peft_settings": {
+            "r": 16,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "target_modules": [
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "o_proj"
+            ],
+            "task_type": "CAUSAL_LM",
+            "modules_to_save": ["embed_tokens", "lm_head"],
+            "name": "LORA"
+        }
+    },
+    "tokenizer_settings": {
+        "tokenizer_path": "/from_s3/model"
+    },
+    "special_tokens_settings": {
+        "bos_token": "<|begin_of_text|>",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|end_of_text|>"
+    },
+    "trainer_settings": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "eval_steps": 250,
+        "save_steps": 250,
+        "per_device_train_batch_size": 2,
+        "per_device_eval_batch_size": 2,
+        "gradient_accumulation_steps": 16,
+        "logging_steps": 1,
+        "learning_rate": 1e-4,
+        "num_train_epochs": 1,
+        "lr_scheduler_type": "cosine",
+        "warmup_ratio": 0.3,
+        "fp16": false,
+        "bf16": true,
+        "optim": "adamw_torch",
+        "save_total_limit": 5,
+        "load_best_model_at_end": false,
+        "deepspeed": "configs/exp/deepspeed/stage2.json",
+        "dispatch_batches": false,
+        "dataloader_num_workers": 16,
+        "dataloader_pin_memory": false
+    },
+    "wandb_settings": {
+        "project_name": "rykov",
+        "run_name": "multimodal",
+        "entity": "rykov"
+    },
+    "log_path": "train_output",
+    "modality_encoder_settings_mapping": {
+        "image": {
+            "modality_encoder_type": "clip",
+            "is_pickle": true,
+            "encoder_path": "/from_s3/clip"
+        },
+        "audio": null
+    },
+    "modality_projector_mapping": {
+        "image": "llava",
+        "audio": null
+    },
+    "modality_projector_initialization_mapping": {
+        "image": null,
+        "audio": null
+    },
+    "cherry_pick_settings": {
+        "generator_transformers_settings": {
+            "num_beams": 1,
+            "max_new_tokens": 64,
+            "repetition_penalty": 1.0,
+            "stop_strings": "<|eot_id|>"
+        },
+        "custom_generation_settings": {
+            "skip_special_tokens": true
+          },
+        "dataset_settings": {
+            "sources": [
+                {
+                    "name": "cherry_pickls",
+                    "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                    "num_samples": 100
+                }
+            ],
+            "prompt_template": {
+                "role_tag_mapping": {
+                    "bot": "assistant",
+                    "user": "user",
+                    "system": "system"
+                },
+                "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+                "suffix_template": "<|eot_id|>"
+            },
+            "dataset_type": "multimodal",
+            "max_tokens_count": 2000,
+            "n_modality_embeddings": 576,
+            "start_modality_token": "<MS>",
+            "end_modality_token": "</MS>",
+            "only_answer_loss": true,
+            "random_cut": true,
+            "modality_token_mapping": {
+                "image": "<img>",
+                "audio": "<audio>"
+            },
+            "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+            "truncate_top": false
+        },
+        "metric_settings": []
+    }
+}
diff --git a/configs/exp/train/multimodal/llava_non_pickled.json b/configs/exp/train/multimodal/llava_non_pickled.json
new file mode 100644
index 0000000..b9ceac4
--- /dev/null
+++ b/configs/exp/train/multimodal/llava_non_pickled.json
@@ -0,0 +1,205 @@
+{
+    "train_dataset_settings": {
+        "sources": [
+            {
+                "name": "train",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/train_chat.jsonl",
+                "num_samples": 100000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+        "n_modality_embeddings": 576,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "val_dataset_settings": {
+        "sources": [
+            {
+                "name": "test",
+                "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                "num_samples": 1000
+            }
+        ],
+        "prompt_template": {
+            "role_tag_mapping": {
+                "bot": "assistant",
+                "user": "user",
+                "system": "system"
+            },
+            "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+            "suffix_template": "<|eot_id|>"
+        },
+        "modality_token_mapping": {
+            "image": "<img>",
+            "audio": "<audio>"
+        },
+        "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+            },
+        "n_modality_embeddings": 576,
+        "start_modality_token": "<MS>",
+        "end_modality_token": "</MS>",
+        "dataset_type": "multimodal",
+        "max_tokens_count": 2000,
+        "only_answer_loss": true,
+        "truncate_top": false
+    },
+    "model_settings": {
+        "model_path": "/from_s3/model",
+        "model_type": "causal",
+        "transformers_settings": {
+        },
+        "model_kwargs": {
+            "attn_implementation": "flash_attention_2"
+        },
+        "embeddings_initialization_strategy": {
+            "<MS>": "bot",
+            "</MS>": "bot",
+            "<img>": "bot",
+            "<audio>": "bot"
+        },
+        "peft_settings": {
+            "r": 16,
+            "lora_alpha": 16,
+            "lora_dropout": 0.05,
+            "target_modules": [
+                "q_proj",
+                "v_proj",
+                "k_proj",
+                "o_proj"
+            ],
+            "task_type": "CAUSAL_LM",
+            "modules_to_save": ["embed_tokens", "lm_head"],
+            "name": "LORA"
+        }
+    },
+    "tokenizer_settings": {
+        "tokenizer_path": "/from_s3/model"
+    },
+    "special_tokens_settings": {
+        "bos_token": "<|begin_of_text|>",
+        "eos_token": "<|im_end|>",
+        "pad_token": "<|end_of_text|>"
+    },
+    "trainer_settings": {
+        "evaluation_strategy": "steps",
+        "save_strategy": "steps",
+        "eval_steps": 100,
+        "save_steps": 100,
+        "per_device_train_batch_size": 1,
+        "per_device_eval_batch_size": 1,
+        "gradient_accumulation_steps": 32,
+        "logging_steps": 1,
+        "learning_rate": 1e-4,
+        "num_train_epochs": 1,
+        "lr_scheduler_type": "cosine",
+        "warmup_ratio": 0.3,
+        "fp16": false,
+        "bf16": true,
+        "optim": "adamw_torch",
+        "save_total_limit": 5,
+        "load_best_model_at_end": false,
+        "deepspeed": "configs/exp/deepspeed/stage2.json",
+        "dispatch_batches": false,
+        "dataloader_num_workers": 0,
+        "dataloader_pin_memory": false
+    },
+    "wandb_settings": {
+        "project_name": "rykov",
+        "run_name": "multimodal",
+        "entity": "rykov"
+    },
+    "log_path": "train_output",
+    "modality_encoder_settings_mapping": {
+        "image": {
+            "modality_encoder_type": "clip",
+            "is_pickle": true,
+            "encoder_path": "/from_s3/clip"
+        },
+        "audio": null
+    },
+    "modality_projector_mapping": {
+        "image": "llava",
+        "audio": null
+    },
+    "modality_projector_initialization_mapping": {
+        "image": null,
+        "audio": null
+    },
+    "cherry_pick_settings": {
+        "generator_transformers_settings": {
+            "num_beams": 1,
+            "max_new_tokens": 64,
+            "repetition_penalty": 1.0,
+            "stop_strings": "<|eot_id|>"
+        },
+        "custom_generation_settings": {
+            "skip_special_tokens": true
+          },
+        "dataset_settings": {
+            "sources": [
+                {
+                    "name": "cherry_pickls",
+                    "records_path": "/from_s3/dataset/llava_next_data_dialogs/test_chat.jsonl",
+                    "num_samples": 25
+                }
+            ],
+            "prompt_template": {
+                "role_tag_mapping": {
+                    "bot": "assistant",
+                    "user": "user",
+                    "system": "system"
+                },
+                "prefix_template": "<|start_header_id|>{role}<|end_header_id|>\n\n",
+                "suffix_template": "<|eot_id|>"
+            },
+            "dataset_type": "multimodal",
+            "max_tokens_count": 2000,
+            "n_modality_embeddings": 576,
+            "start_modality_token": "<MS>",
+            "end_modality_token": "</MS>",
+            "only_answer_loss": true,
+            "random_cut": true,
+            "modality_token_mapping": {
+                "image": "<img>",
+                "audio": "<audio>"
+            },
+            "modality_reader_settings_mapping": {
+                "image": {
+                    "reader_type": "pickle",
+                    "reader_path": null
+                },
+                "audio": null
+        },
+            "truncate_top": false
+        },
+        "metric_settings": []
+    }
+}
diff --git a/configs/utils/preprocess/accelerate_llava_next_data.json b/configs/utils/preprocess/accelerate_llava_next_data.json
new file mode 100644
index 0000000..b52d1c9
--- /dev/null
+++ b/configs/utils/preprocess/accelerate_llava_next_data.json
@@ -0,0 +1,14 @@
+{
+    "reader_settings": {
+        "reader_type": "clip",
+        "reader_path": "/from_s3/clip"
+    },
+    "encoder_settings": {
+        "modality_encoder_type": "clip",
+        "encoder_path": "/from_s3/clip"
+    },
+    "modality": "image",
+    "dataset_path": "/app/llava_next_data/llava_next_data_images",
+    "batch_size": 512,
+    "output_file_path": "/app/output/"
+}
diff --git a/configs/utils/preprocess/coco2014_clip.json b/configs/utils/preprocess/coco2014_clip.json
index 96e765c..b52d1c9 100644
--- a/configs/utils/preprocess/coco2014_clip.json
+++ b/configs/utils/preprocess/coco2014_clip.json
@@ -1,11 +1,14 @@
 {
-    "reader": "clip",
+    "reader_settings": {
+        "reader_type": "clip",
+        "reader_path": "/from_s3/clip"
+    },
     "encoder_settings": {
         "modality_encoder_type": "clip",
-        "encoder_path": "openai/clip-vit-large-patch14"
+        "encoder_path": "/from_s3/clip"
     },
     "modality": "image",
-    "dataset_path": "/app/data/coco2014/images/train2014",
-    "batch_size": 128,
-    "output_file_path": "/output"
+    "dataset_path": "/app/llava_next_data/llava_next_data_images",
+    "batch_size": 512,
+    "output_file_path": "/app/output/"
 }
diff --git a/configs/utils/preprocess/llava_next_data.json b/configs/utils/preprocess/llava_next_data.json
new file mode 100644
index 0000000..b52d1c9
--- /dev/null
+++ b/configs/utils/preprocess/llava_next_data.json
@@ -0,0 +1,14 @@
+{
+    "reader_settings": {
+        "reader_type": "clip",
+        "reader_path": "/from_s3/clip"
+    },
+    "encoder_settings": {
+        "modality_encoder_type": "clip",
+        "encoder_path": "/from_s3/clip"
+    },
+    "modality": "image",
+    "dataset_path": "/app/llava_next_data/llava_next_data_images",
+    "batch_size": 512,
+    "output_file_path": "/app/output/"
+}
diff --git a/turbo_alignment/common/data/multimodal/common.py b/turbo_alignment/common/data/multimodal/common.py
index cb17769..72449d8 100644
--- a/turbo_alignment/common/data/multimodal/common.py
+++ b/turbo_alignment/common/data/multimodal/common.py
@@ -24,4 +24,5 @@ def _get_h5_file(path: Path) -> Path:
     def read(self, path: str) -> torch.Tensor:
         h5_file = self._get_h5_file(Path(path).parent)
         with h5py.File(h5_file, 'r') as f:
-            return torch.tensor(f[Path(path).name])
+            output = torch.tensor(f[Path(path).name])
+        return output
diff --git a/turbo_alignment/dataset/multimodal/collators.py b/turbo_alignment/dataset/multimodal/collators.py
index 28ccc4e..fa39839 100644
--- a/turbo_alignment/dataset/multimodal/collators.py
+++ b/turbo_alignment/dataset/multimodal/collators.py
@@ -7,7 +7,7 @@ class DataCollatorWithModalityInputs(DataCollatorForTokenClassification):
     def torch_call(self, features):
         label_name = 'label' if 'label' in features[0].keys() else 'labels'
         labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
-
+        # print("😍"*10, features[0].keys())
         if 'modality_inputs' in features[0].keys():
             modality_inputs = [feature['modality_inputs'] for feature in features]
         else: