Skip to content

Commit f488ab6

Browse files
sywangyisywangyi
authored
Add mllama support (#1419)
Signed-off-by: Wang, Yi A <yi.a.wang@intel.com>
1 parent 7fe7bf8 commit f488ab6

14 files changed

Lines changed: 1396 additions & 38 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@ The following model architectures, tasks and device distributions have been vali
235235
| VideoMAE | | <div style="text-align:left"><li>Single card</li></div> | <li>[Video classification](https://github.com/huggingface/optimum-habana/tree/main/examples/video-classification)</li> |
236236
| TableTransformer | | <div style="text-align:left"><li>Single card</li></div> | <li>[table object detection](https://github.com/huggingface/optimum-habana/tree/main/examples/table-detection) </li> |
237237
| DETR | | <div style="text-align:left"><li>Single card</li></div> | <li>[object detection](https://github.com/huggingface/optimum-habana/tree/main/examples/object-detection)</li> |
238+
| Mllama | <div style="text-align:left"><li>LoRA</li></div> | :heavy_check_mark: | <li>[image to text](https://github.com/huggingface/optimum-habana/tree/main/examples/image-to-text)</li> |
238239
239240
</div>
240241

‎docs/source/index.mdx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ In the tables below, ✅ means single-card, multi-card and DeepSpeed have all be
8181
| VideoMAE | | <div style="text-align:left"><li>Single card</li></div> | <li>[Video classification](https://github.com/huggingface/optimum-habana/tree/main/examples/video-classification)</li> |
8282
| TableTransformer | | <div style="text-align:left"><li>Single card</li></div> | <li>[table object detection](https://github.com/huggingface/optimum-habana/tree/main/examples/table-detection)</li> |
8383
| DETR | | <div style="text-align:left"><li>Single card</li></div> | <li>[object detection](https://github.com/huggingface/optimum-habana/tree/main/examples/object-detection)</li> |
84+
| Mllama | <div style="text-align:left"><li>LoRA</li></div> |✅ | <li>[image to text](https://github.com/huggingface/optimum-habana/tree/main/examples/image-to-text)</li> |
8485

8586
- Diffusers
8687

‎examples/image-to-text/README.md‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ Models that have been validated:
3131
- [llava-hf/llava-v1.6-34b-hf](https://huggingface.co/llava-hf/llava-v1.6-34b-hf)
3232
- [llava-hf/llama3-llava-next-8b-hf](https://huggingface.co/llava-hf/llama3-llava-next-8b-hf)
3333
- [HuggingFaceM4/idefics2-8b](https://huggingface.co/HuggingFaceM4/idefics2-8b)
34+
- [meta-llama/Llama-3.2-11B-Vision-Instruct](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct)
3435

3536
### Inference with BF16
3637

@@ -102,6 +103,15 @@ python3 run_pipeline.py \
102103
--bf16
103104
```
104105

106+
To run mllama inference, use the following command:
107+
108+
```bash
109+
python3 run_pipeline.py \
110+
--model_name_or_path meta-llama/Llama-3.2-11B-Vision-Instruct \
111+
--use_hpu_graphs \
112+
--bf16
113+
```
114+
105115
### Inference with FP8
106116
Inference for Llava-1.5-7b, Llava-1.5-13b, Llava-v1.6-mistral-7b and Llava-v1.6-vicuna-13b in FP8 precision are enabled using [Intel Neural Compressor (INC)](https://docs.habana.ai/en/latest/PyTorch/Inference_on_PyTorch/Inference_Using_FP8.html), which provides model measurement and quantization capabilities in PyTorch.
107117

@@ -286,6 +296,75 @@ python3 ../gaudi_spawn.py \
286296
--lora_target_modules '".*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$"'
287297
```
288298

299+
Here are single-/multi-device command examples for meta-llama/Llama-3.2-11B-Vision-Instruct.
300+
301+
```bash
302+
python3 run_image2text_lora_finetune.py \
303+
--model_name_or_path meta-llama/Llama-3.2-11B-Vision-Instruct \
304+
--dataset_name nielsr/docvqa_1200_examples \
305+
--bf16 True \
306+
--output_dir ./model_lora_llama \
307+
--num_train_epochs 2 \
308+
--per_device_train_batch_size 2 \
309+
--per_device_eval_batch_size 2 \
310+
--gradient_accumulation_steps 8 \
311+
--weight_decay 0.01 \
312+
--logging_steps 25 \
313+
--eval_strategy "no" \
314+
--save_strategy "no" \
315+
--learning_rate 5e-5 \
316+
--warmup_steps 50 \
317+
--lr_scheduler_type "constant" \
318+
--input_column_names 'image' 'query' \
319+
--output_column_names 'answers' \
320+
--remove_unused_columns False \
321+
--do_train \
322+
--do_eval \
323+
--use_habana \
324+
--use_lazy_mode \
325+
--lora_rank=8 \
326+
--lora_alpha=8 \
327+
--lora_dropout=0.1 \
328+
--low_cpu_mem_usage True \
329+
--max_seq_length=512 \
330+
--use_hpu_graphs_for_inference True \
331+
--lora_target_modules ".*(language_model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$"
332+
```
333+
334+
```bash
335+
python3 ../gaudi_spawn.py \
336+
--world_size 8 --use_mpi run_image2text_lora_finetune.py \
337+
--model_name_or_path meta-llama/Llama-3.2-11B-Vision-Instruct \
338+
--dataset_name nielsr/docvqa_1200_examples \
339+
--bf16 True \
340+
--output_dir ./model_lora_llama \
341+
--num_train_epochs 2 \
342+
--per_device_train_batch_size 2 \
343+
--per_device_eval_batch_size 2 \
344+
--gradient_accumulation_steps 8 \
345+
--weight_decay 0.01 \
346+
--logging_steps 25 \
347+
--eval_strategy "no" \
348+
--save_strategy "no" \
349+
--learning_rate 5e-5 \
350+
--warmup_steps 50 \
351+
--lr_scheduler_type "constant" \
352+
--input_column_names 'image' 'query' \
353+
--output_column_names 'answers' \
354+
--remove_unused_columns False \
355+
--do_train \
356+
--do_eval \
357+
--use_habana \
358+
--use_lazy_mode \
359+
--lora_rank=8 \
360+
--lora_alpha=8 \
361+
--lora_dropout=0.1 \
362+
--low_cpu_mem_usage True \
363+
--max_seq_length=512 \
364+
--use_hpu_graphs_for_inference True \
365+
--lora_target_modules '".*(language_model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$"'
366+
```
367+
289368
## Multi-HPU inference
290369

291370
To enable multi-card inference, you must set the environment variable `PT_HPU_ENABLE_LAZY_COLLECTIVES=true`,

‎examples/image-to-text/run_image2text_lora_finetune.py‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -251,11 +251,9 @@ class FinetuneArguments:
251251

252252

253253
class MyDataCollator:
254-
def __init__(self, processor, max_seq_length):
254+
def __init__(self, processor, max_seq_length, image_token_id):
255255
self.processor = processor
256-
self.image_token_id = processor.tokenizer.additional_special_tokens_ids[
257-
processor.tokenizer.additional_special_tokens.index("<image>")
258-
]
256+
self.image_token_id = image_token_id
259257
self.max_seq_length = max_seq_length
260258

261259
def __call__(self, examples):
@@ -458,8 +456,15 @@ def main():
458456
if col not in (data_args.input_column_names + data_args.output_column_names)
459457
]
460458
)
461-
462-
data_collator = MyDataCollator(processor, max_seq_length=data_args.max_seq_length)
459+
if hasattr(config, "image_token_id"):
460+
# idefics
461+
image_token_id = config.image_token_id
462+
elif hasattr(config, "image_token_index"):
463+
# mllama
464+
image_token_id = config.image_token_index
465+
else:
466+
raise ValueError("Please provide value for image_token_id")
467+
data_collator = MyDataCollator(processor, max_seq_length=data_args.max_seq_length, image_token_id=image_token_id)
463468

464469
gaudi_config = GaudiConfig()
465470
gaudi_config.use_fused_adam = True

‎examples/image-to-text/run_pipeline.py‎

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import PIL.Image
2424
import requests
2525
import torch
26-
from transformers import AutoConfig, AutoProcessor, pipeline
26+
from transformers import AutoConfig, AutoModelForVision2Seq, AutoProcessor, pipeline
2727

2828
from optimum.habana.transformers.modeling_utils import adapt_transformers_to_gaudi
2929

@@ -185,14 +185,13 @@ def main():
185185
adapt_transformers_to_gaudi()
186186

187187
model_type = AutoConfig.from_pretrained(args.model_name_or_path).model_type
188-
if args.image_path is None and model_type in ["llava", "idefics2"]:
188+
if args.image_path is None and model_type in ["llava", "idefics2", "mllama"]:
189189
args.image_path = ["https://llava-vl.github.io/static/images/view.jpg"]
190190
elif args.image_path is None and model_type == "llava_next":
191191
args.image_path = [
192192
"https://github.com/haotian-liu/LLaVA/blob/1a91fc274d7c35a9b50b3cb29c4247ae5837ce39/images/llava_v1_5_radar.jpg?raw=true"
193193
]
194-
195-
if args.prompt is None and model_type in ["llava", "idefics2", "llava_next"]:
194+
if args.prompt is None and model_type in ["llava", "idefics2", "llava_next", "mllama"]:
196195
processor = AutoProcessor.from_pretrained(args.model_name_or_path)
197196
conversation = [
198197
{
@@ -231,17 +230,31 @@ def main():
231230

232231
htcore.hpu_set_env()
233232

234-
generator = pipeline(
235-
"image-to-text",
236-
model=args.model_name_or_path,
237-
torch_dtype=model_dtype,
238-
device="hpu",
239-
)
240-
241233
if args.world_size > 1:
242-
generator.model = initialize_distributed_model(args, generator.model, logger, model_dtype)
243-
234+
import deepspeed
235+
236+
with deepspeed.OnDevice(dtype=model_dtype, device="cpu"):
237+
model = AutoModelForVision2Seq.from_pretrained(args.model_name_or_path, torch_dtype=model_dtype)
238+
if model_type == "mllama":
239+
model.language_model = initialize_distributed_model(args, model.language_model, logger, model_dtype)
240+
else:
241+
model = initialize_distributed_model(args, model, logger, model_dtype)
242+
generator = pipeline(
243+
"image-to-text",
244+
model=model,
245+
config=args.model_name_or_path,
246+
tokenizer=args.model_name_or_path,
247+
image_processor=args.model_name_or_path,
248+
torch_dtype=model_dtype,
249+
device="hpu",
250+
)
244251
else:
252+
generator = pipeline(
253+
"image-to-text",
254+
model=args.model_name_or_path,
255+
torch_dtype=model_dtype,
256+
device="hpu",
257+
)
245258
if args.use_hpu_graphs:
246259
from habana_frameworks.torch.hpu import wrap_in_hpu_graph
247260

@@ -263,7 +276,7 @@ def main():
263276
htcore.hpu_initialize(generator.model)
264277

265278
# delete once pipeline integrate AutoProcessor as preprocess engine
266-
if model_type in ["idefics2"]:
279+
if model_type in ["idefics2", "mllama"]:
267280
from transformers.image_utils import load_image
268281

269282
def preprocess(self, image, prompt=None, timeout=None):

‎optimum/habana/transformers/generation/utils.py‎

Lines changed: 36 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@
109109
"xglm",
110110
"whisper",
111111
"idefics2",
112+
"mllama",
112113
]
113114

114115

@@ -330,11 +331,13 @@ def _expand_dict_for_generation(dict_to_expand):
330331

331332
def _pad_past_key_values(self, model_kwargs):
332333
pad_amount = model_kwargs.get("kv_cache_pad_len", 0)
334+
kv_cache_len = model_kwargs.get("kv_cache_len", 0)
333335
if model_kwargs["past_key_values"]:
334336
if model_kwargs.get("mqa_model", False):
335337
for i in range(len(model_kwargs["past_key_values"])): # layer
336-
if torch.is_tensor(
337-
model_kwargs["past_key_values"][i]
338+
if (
339+
torch.is_tensor(model_kwargs["past_key_values"][i])
340+
and model_kwargs["past_key_values"][i].shape[-2] == kv_cache_len - pad_amount
338341
): # tensor(batch_size, kv_cache_len, n_heads * head_dim * 2) k and v stacked
339342
model_kwargs["past_key_values"][i] = torch.nn.functional.pad(
340343
model_kwargs["past_key_values"][i], (0, 0, 0, pad_amount)
@@ -344,8 +347,9 @@ def _pad_past_key_values(self, model_kwargs):
344347
else:
345348
for i in range(len(model_kwargs["past_key_values"])): # layer
346349
for j in range(len(model_kwargs["past_key_values"][i])): # k or v
347-
if torch.is_tensor(
348-
model_kwargs["past_key_values"][i][j]
350+
if (
351+
torch.is_tensor(model_kwargs["past_key_values"][i][j])
352+
and model_kwargs["past_key_values"][i][j].shape[-2] == kv_cache_len - pad_amount
349353
): # tensor(batch_size, n_heads, kv_cache_len, head_dim)
350354
model_kwargs["past_key_values"][i][j] = torch.nn.functional.pad(
351355
model_kwargs["past_key_values"][i][j], (0, 0, 0, pad_amount)
@@ -461,6 +465,14 @@ def update_model_kwargs_for_bucketing(
461465
)
462466
else:
463467
assert False, "Not tested for cases where attn_mask isnt passed"
468+
469+
if model_kwargs.get("cross_attention_mask") is not None:
470+
model_kwargs["cross_attention_mask"] = torch.nn.functional.pad(
471+
model_kwargs["cross_attention_mask"],
472+
(0, 0, 0, 0, 0, pad_amount),
473+
value=0,
474+
)
475+
464476
if reduce_recompile and params["passnum"] == 0:
465477
position_ids_cpu = model_kwargs["attention_mask"].long().cumsum(-1) - 1
466478
position_ids_cpu.masked_fill_(model_kwargs["attention_mask"] == 0, 1)
@@ -503,14 +515,20 @@ def create_pad_arg(pad_amount, i, j):
503515
# This is a necessary (but not sufficient) condition: what ever dimension we are padding, should be a multiple of bucket_size
504516
# This check is added in case we get a new model with a new kv-cache structure, and we attempt to pad some wrong dimension
505517
# in peft case, if there's virtual token. the model_kwargs["past_key_values"][i][j].shape[-(len(pad_tuple) // 2)] % bucket_size == num_virtual_token, no need of assert, the pad length of past_key_value should be aligned with input id and attention_mask
506-
num_virtual_tokens = model_kwargs.get("num_virtual_tokens", 0)
507-
assert (
508-
model_kwargs["past_key_values"][i][j].shape[-(len(pad_tuple) // 2)] % bucket_size
509-
== num_virtual_tokens
510-
)
511-
tmp_lst[j] = torch.nn.functional.pad(
512-
model_kwargs["past_key_values"][i][j], pad_tuple, value=pad_token_id
513-
)
518+
if (
519+
model_kwargs["past_key_values"][i][j].shape[-(len(pad_tuple) // 2)]
520+
== params["allocated_space"] - pad_amount
521+
):
522+
num_virtual_tokens = model_kwargs.get("num_virtual_tokens", 0)
523+
assert (
524+
model_kwargs["past_key_values"][i][j].shape[-(len(pad_tuple) // 2)] % bucket_size
525+
== num_virtual_tokens
526+
)
527+
tmp_lst[j] = torch.nn.functional.pad(
528+
model_kwargs["past_key_values"][i][j], pad_tuple, value=pad_token_id
529+
)
530+
else:
531+
tmp_lst[j] = model_kwargs["past_key_values"][i][j]
514532
new_kv[i] = tuple(tmp_lst)
515533
model_kwargs["past_key_values"] = tuple(new_kv)
516534

@@ -1110,6 +1128,12 @@ def generate(
11101128
(0, generation_config.max_new_tokens),
11111129
value=0,
11121130
)
1131+
if model_kwargs.get("cross_attention_mask") is not None:
1132+
model_kwargs["cross_attention_mask"] = torch.nn.functional.pad(
1133+
model_kwargs["cross_attention_mask"],
1134+
(0, 0, 0, 0, 0, generation_config.max_new_tokens),
1135+
value=0,
1136+
)
11131137
else:
11141138
assert generation_config.bucket_size <= 0, "Untested path for bucket>0"
11151139
if model_kwargs.get("decoder_input_ids", None) is None:

‎optimum/habana/transformers/modeling_utils.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,14 @@
8888
GaudiMixtralDecoderLayer,
8989
GaudiMixtralForCausalLM,
9090
GaudiMixtralModel,
91+
GaudiMllamaCrossAttentionDecoderLayer,
92+
GaudiMllamaForCausalLM,
93+
GaudiMllamaForConditionalGeneration,
94+
GaudiMllamaSelfAttentionDecoderLayer,
95+
GaudiMllamaTextCrossAttention,
96+
GaudiMllamaTextModel,
97+
GaudiMllamaTextSelfAttention,
98+
GaudiMllamaVisionModel,
9199
GaudiMptAttention,
92100
GaudiMptBlock,
93101
GaudiMptForCausalLM,
@@ -622,6 +630,16 @@ def adapt_transformers_to_gaudi():
622630
transformers.models.whisper.modeling_whisper.WhisperForConditionalGeneration = GaudiWhisperForConditionalGeneration
623631
transformers.models.whisper.modeling_whisper.WHISPER_ATTENTION_CLASSES = GAUDI_WHISPER_ATTENTION_CLASSES
624632

633+
# Optimization for mllama on Gaudi
634+
transformers.models.mllama.modeling_mllama.MllamaSelfAttentionDecoderLayer = GaudiMllamaSelfAttentionDecoderLayer
635+
transformers.models.mllama.modeling_mllama.MllamaCrossAttentionDecoderLayer = GaudiMllamaCrossAttentionDecoderLayer
636+
transformers.models.mllama.modeling_mllama.MllamaForCausalLM = GaudiMllamaForCausalLM
637+
transformers.models.mllama.modeling_mllama.MllamaTextSelfAttention = GaudiMllamaTextSelfAttention
638+
transformers.models.mllama.modeling_mllama.MllamaTextCrossAttention = GaudiMllamaTextCrossAttention
639+
transformers.models.mllama.modeling_mllama.MllamaForConditionalGeneration = GaudiMllamaForConditionalGeneration
640+
transformers.models.mllama.modeling_mllama.MllamaTextModel = GaudiMllamaTextModel
641+
transformers.models.mllama.modeling_mllama.MllamaVisionModel = GaudiMllamaVisionModel
642+
625643
transformers.AutoConfig.register("deci", DeciLMConfig)
626644
transformers.AutoModelForCausalLM.register(DeciLMConfig, DeciLMForCausalLM)
627645

‎optimum/habana/transformers/models/__init__.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,16 @@
150150
gaudi_mixtral_block_sparse_moe_forward,
151151
gaudi_mixtral_rmsnorm_forward,
152152
)
153+
from .mllama import (
154+
GaudiMllamaCrossAttentionDecoderLayer,
155+
GaudiMllamaForCausalLM,
156+
GaudiMllamaForConditionalGeneration,
157+
GaudiMllamaSelfAttentionDecoderLayer,
158+
GaudiMllamaTextCrossAttention,
159+
GaudiMllamaTextModel,
160+
GaudiMllamaTextSelfAttention,
161+
GaudiMllamaVisionModel,
162+
)
153163
from .modeling_all_models import (
154164
gaudi_check_and_enable_sdpa,
155165
gaudi_conv1d_forward,

‎optimum/habana/transformers/models/mllama/__init__.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
from .modeling_mllama import (
2+
GaudiMllamaCrossAttentionDecoderLayer,
3+
GaudiMllamaForCausalLM,
4+
GaudiMllamaForConditionalGeneration,
5+
GaudiMllamaSelfAttentionDecoderLayer,
6+
GaudiMllamaTextCrossAttention,
7+
GaudiMllamaTextModel,
8+
GaudiMllamaTextSelfAttention,
9+
GaudiMllamaVisionModel,
10+
)

0 commit comments

Comments
 (0)