-
Notifications
You must be signed in to change notification settings - Fork 25.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add trainer integration test for llava to ensure accelerate autocasting works correctly #30489
base: main
Are you sure you want to change the base?
Changes from all commits
5aa7002
3c669fa
a0fa2ae
bc31529
3c7e7e1
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,116 @@ | ||
import gc | ||
|
||
import requests | ||
from datasets import Dataset | ||
|
||
from transformers import ( | ||
AutoProcessor, | ||
BitsAndBytesConfig, | ||
DataCollatorForLanguageModeling, | ||
LlavaForConditionalGeneration, | ||
Trainer, | ||
TrainingArguments, | ||
is_torch_available, | ||
is_vision_available, | ||
) | ||
from transformers.testing_utils import TestCasePlus, require_bitsandbytes, require_peft, require_torch, slow | ||
|
||
|
||
if is_vision_available(): | ||
from PIL import Image | ||
|
||
if is_torch_available(): | ||
import torch | ||
else: | ||
is_torch_greater_or_equal_than_2_0 = False | ||
|
||
|
||
# Integration test for confirming autocast with trainer and accelerate works | ||
# correctly. Confirms type error found | ||
# https://github.com/huggingface/transformers/pull/29721 in is fixed | ||
@require_torch | ||
class LlavaForConditionalGenerationIntegrationTest(TestCasePlus): | ||
def setUp(self): | ||
super().setUp() | ||
self.processor = AutoProcessor.from_pretrained( | ||
"llava-hf/bakLlava-v1-hf", padding_side="left", truncation_side="right" | ||
) | ||
|
||
def tearDown(self): | ||
gc.collect() | ||
torch.cuda.empty_cache() | ||
|
||
@slow | ||
@require_bitsandbytes | ||
@require_peft | ||
def test_model_trainer_integration_test(self): | ||
from peft import LoraConfig, PeftModelForCausalLM | ||
|
||
def image_prompt_generator(): | ||
prompts = [ | ||
"USER: <image>\nWhat are the things I should be cautious about when I visit this place? What should I bring with me?\nASSISTANT:", | ||
"USER: <image>\nWhat is this?\nASSISTANT:", | ||
] | ||
image_urls = [ | ||
"https://llava-vl.github.io/static/images/view.jpg", | ||
"http://images.cocodataset.org/val2017/000000039769.jpg", | ||
] | ||
|
||
for prompt, image_url in zip(prompts, image_urls): | ||
image = Image.open(requests.get(image_url, stream=True).raw) | ||
yield {"image": image, "prompt": prompt} | ||
|
||
def process_image_prompt(data): | ||
processed = self.processor( | ||
data["prompt"], images=data["image"], return_tensors="pt", padding=True, max_length=512 | ||
) | ||
return { | ||
"input_ids": processed["input_ids"].squeeze(), | ||
"attention_mask": processed["attention_mask"].squeeze(), | ||
"pixel_values": processed["pixel_values"].squeeze(), | ||
} | ||
|
||
train_dataset = Dataset.from_generator(image_prompt_generator).map(process_image_prompt) | ||
bits_and_bytes_config = BitsAndBytesConfig( | ||
load_in_4bit=True, | ||
) | ||
model = LlavaForConditionalGeneration.from_pretrained( | ||
"llava-hf/bakLlava-v1-hf", quantization_config=bits_and_bytes_config | ||
) | ||
peft_config = LoraConfig( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I am somewhat unclear on where we test the slow tests but I assume there is some limit on memory so I tried to give a reasonable LORA for this test. If you think there is a simpler or more idiomatic way to do this test please let me know. |
||
r=16, | ||
lora_alpha=16, | ||
bias="none", | ||
task_type="CAUSAL_LM", | ||
lora_dropout=0.0, | ||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], | ||
) | ||
model = PeftModelForCausalLM(model, peft_config, adapter_name="lora_default") | ||
data_collator = DataCollatorForLanguageModeling(self.processor.tokenizer, mlm=False) | ||
|
||
output_dir = self.get_auto_remove_tmp_dir() | ||
trainer = Trainer( | ||
model=model, | ||
train_dataset=train_dataset, | ||
tokenizer=self.processor.tokenizer, | ||
args=TrainingArguments(output_dir, fp16=True, learning_rate=2e-5, num_train_epochs=1), | ||
data_collator=data_collator, | ||
) | ||
trainer.train() | ||
|
||
prompts = [ | ||
"USER: <image>\nWhat are the things I should be cautious about when I visit this place? What should I bring with me?\nASSISTANT:", | ||
"USER: <image>\nWhat is this?\nASSISTANT:", | ||
] | ||
image1 = Image.open(requests.get("https://llava-vl.github.io/static/images/view.jpg", stream=True).raw) | ||
image2 = Image.open(requests.get("http://images.cocodataset.org/val2017/000000039769.jpg", stream=True).raw) | ||
|
||
inputs = self.processor(prompts, images=[image1, image2], return_tensors="pt", padding=True) | ||
|
||
output = model(**inputs) | ||
expected_slice = torch.tensor( | ||
[[-3.5664, -3.5625, -0.4309], [-5.8242, -5.6914, -1.3242], [-5.4805, -5.9375, 1.1465]], | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I saw this pattern in the llava next tests. These came from the use of this model before training. Not quite sure if this is correct so please let me know if there is something else we want to test. Perhaps instead we want the trained model before applying the downcasting change? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Also note these do not yet pass |
||
dtype=torch.float32, | ||
) | ||
|
||
assert torch.allclose(output["logits"][0, :3, :3], expected_slice, atol=1e-3) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not entirely sure this is the simplest or most idiomatic way to create this test dataset so please let me know if there is a better way.