Skip to content

[Draft] Add autocast to prediction_step for SFTTrainer#2310

Closed
pdufour wants to merge 19 commits intohuggingface:mainfrom
pdufour:train_predict
Closed

[Draft] Add autocast to prediction_step for SFTTrainer#2310
pdufour wants to merge 19 commits intohuggingface:mainfrom
pdufour:train_predict

Conversation

@pdufour
Copy link
Copy Markdown

@pdufour pdufour commented Nov 3, 2024 •
edited
Loading

What does this PR do?

Note: not ready for review until huggingface/transformers#32346 lands.

There is a PR to add predict_with_generate to the base Trainer class https://github.com/huggingface/transformers/pull/32346/files. However without any changes to SFTrainer, it throws error if you use it and have mixed precision turned on.

You can see a simliar bug (and fix) happened here - #1203

This pull request aims to make SFTrainer compatible with the soon-to-land predict_with_generate changes.

Reproducible Example:

import torch
from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer
from datasets import Dataset
from peft import LoraConfig
from PIL import Image
from io import BytesIO

model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)
model = AutoModelForVision2Seq.from_pretrained(model_id, device_map="auto", torch_dtype=torch.bfloat16, quantization_config=bnb_config)

processor = AutoProcessor.from_pretrained(model_id)
tokenizer = processor.tokenizer

peft_config = LoraConfig(target_modules=["k_proj", "o_proj", "q_proj", "v_proj"], task_type="CAUSAL_LM")

args = SFTConfig(
    max_steps=2, output_dir="/tmp/test-model", gradient_checkpointing=True, bf16=True, bf16_full_eval=True,
    logging_steps=1, push_to_hub=False, report_to="wandb", eval_strategy="steps", eval_steps=2,
    generation_config={"num_beams": 1, "max_length": 1024}, predict_with_generate=True,
    gradient_checkpointing_kwargs={"use_reentrant": False}, dataset_kwargs={"skip_prepare_dataset": True},
    remove_unused_columns=False
)

# Example of structured data without URLs for images
data = [
    {
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "Generate code for a web page that looks exactly like this."},
                    {"type": "image", "image": "placeholder"}
                ]
            },
            {
                "role": "assistant",
                "content": [
                    {"type": "text", "text": "<html><body>Sample HTML content for the web page...</body></html>"}
                ]
            }
        ]
    }
]

dataset = Dataset.from_list(data)

def generate_placeholder_image():
    placeholder_image = Image.new("RGB", (256, 256), color="gray")
    return placeholder_image

def collate_fn(examples, mode):
    texts = [processor.apply_chat_template(example["messages"], tokenize=False).strip() for example in examples]
    texts_eval = [processor.apply_chat_template([example["messages"][0]], add_generation_prompt=True).strip() for example in examples]

    images = [[generate_placeholder_image()] for _ in examples]
    
    processor.tokenizer.padding_side = "right"
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
    
    if mode == 'test':
        processor.tokenizer.padding_side = "left"
        batch_eval = processor(text=texts_eval, images=images, return_tensors="pt", padding=True)
        batch["generation_input_ids"] = batch_eval["input_ids"]
        batch["generation_attention_mask"] = batch_eval["attention_mask"]

    labels = batch["input_ids"].clone()
    labels[labels == tokenizer.pad_token_id] = -100
    image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token)
    labels[labels == image_token_id] = -100
    batch["labels"] = labels
    
    return batch

trainer = SFTTrainer(
    model=model, args=args, train_dataset=dataset, data_collator=lambda x: collate_fn(x, mode='train'), eval_dataset=dataset, 
    dataset_text_field="messages", tokenizer=tokenizer, peft_config=peft_config,
    compute_metrics = lambda: dict(),
)

trainer.eval_data_collator = lambda examples: collate_fn(examples, 'test')

trainer.train()
print('Done')

Test Plan

Test fail case

  1. Install transformers branch with train_predict flag pip install git+https://github.com/pdufour/transformers.git@train_predict
  2. Edit the above code and comment out the generation config, predict_with_generate=True, arg, and print(f"bf16 casting status: {trainer._peft_has_been_casted_to_bf16}") line
  3. Run the above code
  4. See the error
image

Test success case

  1. Install trl branch with this fix (this pr)pip install git+https://github.com/pdufour/trl.git@train_predict
  2. Run the code
  3. See no error

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a GitHub issue? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@pdufour pdufour changed the title Add autocast to prediction_step Add autocast to prediction_step for SFTTrainer Nov 3, 2024
@pdufour pdufour changed the title Add autocast to prediction_step for SFTTrainer [Draft] Add autocast to prediction_step for SFTTrainer Nov 3, 2024
@pdufour pdufour mentioned this pull request Nov 3, 2024
Open
@qgallouedec
Copy link
Copy Markdown
Member

qgallouedec commented Jul 22, 2025

I'm closing this PR because no recent activity. Please feel free to reopen it or request a reopening if necessary.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@pdufour @qgallouedec