Vision Datasets

[kylesayrs]

Background

Lifecycle of a dataset

Action	Data form	Purpose	Where
Load dataset	String, dataset	Potentially load from HF/DVC	TextGenerationDataset
Preprocess (apply chat template)	Raw dataset	Format data for processor	TextGenerationDataset.preprocess
Rename columns	Preprocessed dataset + prompt key	Support datasets with text columns not named "text"	TextGenerationDataset.rename_columns
Filter tokenizer args	Preprocessed dataset + prompt key	Remove columns not used by tokenizer (automates data_args.remove_columns	TextGenerationDataset.filter_tokenizer_args
Tokenize (process)	Dataset of processor kwargs + prompt key	Tokenize into model kwargs	TextGenerationDataset.tokenize
Postprocess	Dataset of model kwargs	Potentially concatenate data, add labels for finetuning	TextGenerationDataset.group_text and/or TextGenerationDataset.add_labels
Apply dataloader	Dataset of model kwargs	Batch model kwargs for calibration/training	format_calibration_data or Trainer.get_X_dataloader
Forward pass	Model kwargs	Train or calibrate model	PretrainedModel.forward

• The user can pass a dataset in any one of the above intermediate forms, and TextGenerationDataset will skip all previous steps and only perform the necessary subsequent steps. Typically this means passing a dataset name, raw dataset, or tokenized dataset.

Purpose

• Clean up/ clarify dataset lifecycle
• Support vision processors, which means supporting arbitrary arguments as input
  • Previously we relied on the user specifying the "text column" which was the only input into the tokenizer
  • Vision processors take text as input, but also many other arguments such as pixel_values, aspect_ratio, ect.
• No longer require users to specify remove_columns, instead infer this information from the tokenizer call signature
• Allow users to pass their own custom data_collator function, which is necessary for many vision datasets
• Allow preprocessing_func to be used with any dataset

Changes

• Refactor TextGenerationDataset class to clarify dataset lifecycle
• Add Flickr30K dataset
• Depreciate remove_columns argument
• Add data_collator argument

Prerequisites

• Replace tokenizer with processor #955
  • Remove unused sparseml.export utilities #950

Sample Debug Logs

2024-12-20 20:59:00.115 | DEBUG    | llmcompressor.transformers.finetune.data.base:load_dataset:179 - Loading dataset lmms-lab/flickr30k
2024-12-20 20:59:01.073 | DEBUG    | llmcompressor.transformers.finetune.data.base:__call__:91 - Raw dataset: ['image', 'caption', 'sentids', 'img_id', 'filename']
2024-12-20 20:59:01.171 | DEBUG    | llmcompressor.transformers.finetune.data.base:__call__:102 - Dataset after preprocessing: ['image', 'caption', 'sentids', 'img_id', 'filename', 'text', 'images']
2024-12-20 20:59:01.171 | DEBUG    | llmcompressor.transformers.finetune.data.base:__call__:106 - Dataset after column renaming: ['image', 'caption', 'sentids', 'img_id', 'filename', 'text', 'images']
2024-12-20 20:59:01.171 | DEBUG    | llmcompressor.transformers.finetune.data.base:filter_tokenizer_args:232 - Found processor args `{'images', 'text', 'videos'}`. Removing all other columns
2024-12-20 20:59:01.172 | DEBUG    | llmcompressor.transformers.finetune.data.base:__call__:111 - Tokenizer args after filtering: ['text', 'images']
Tokenizing: 100%|???????????????????????????????????????????????????????????????????????????????????????????????????????????????????????????????| 3/3 [00:00<00:00, 61.32 examples/s]
2024-12-20 20:59:01.289 | DEBUG    | llmcompressor.transformers.finetune.data.base:__call__:125 - Model kwargs after tokenizing: ['input_ids', 'attention_mask', 'pixel_values', 'image_grid_thw']
2024-12-20 20:59:01.289 | DEBUG    | llmcompressor.transformers.finetune.data.base:__call__:155 - Model kwargs after postprocessing: ['input_ids', 'attention_mask', 'pixel_values', 'image_grid_thw']

Testing

• Able to run custom dataset examples/quantization_w4a16/llama3_example.py
• Use the following scripts to verify that the flickr vision dataset loads and can be inputted into the model (although only the first layer will quantize due to missing changes from VLM Support via GPTQ Hooks and Data Pipelines #914)
• Nightly ✅

mllama.py

import os

import torch
from transformers import AutoProcessor, MllamaForConditionalGeneration

from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.transformers import oneshot

# Load model.
model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
model = MllamaForConditionalGeneration.from_pretrained(
    model_id, device_map="auto", torch_dtype="auto"
)
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

# Oneshot arguments
DATASET_ID = "flickr30k"
DATASET_SPLIT = "test[:512]"
NUM_CALIBRATION_SAMPLES = 1
MAX_SEQUENCE_LENGTH = 2048


# TODO: define real collators in utils
def data_collator(batch):
    assert len(batch) == 1
    return {
        "input_ids": torch.LongTensor(batch[0]["input_ids"]),
        "attention_mask": torch.tensor(batch[0]["attention_mask"]),
        "pixel_values": torch.tensor(batch[0]["pixel_values"]),
        "aspect_ratio_ids": torch.tensor(batch[0]["aspect_ratio_ids"]),
        "aspect_ratio_mask": torch.tensor(batch[0]["aspect_ratio_mask"]),
        "cross_attention_mask": torch.tensor(batch[0]["cross_attention_mask"]),
    }


# Recipe
recipe = [
    # SmoothQuantModifier(smoothing_strength=0.8, ignore=ignore),
    GPTQModifier(
        targets="Linear",
        scheme="W8A8",
        ignore=["re:.*lm_head", "re:multi_modal_projector.*", "re:vision_model.*"],
        update_size=NUM_CALIBRATION_SAMPLES,
    ),
]

# Perform oneshot
save_name = model_id.split("/")[1] + "-W8A8"
save_path = os.path.join("./my_test/", save_name)
print("Starting quantization")
oneshot(
    model=model,
    tokenizer=model_id,
    dataset=DATASET_ID,
    splits=DATASET_SPLIT,
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    trust_remote_code_model=True,
    output_dir=save_path,
    data_collator=data_collator,
)

processor.save_pretrained(save_path)

# Confirm generations of the quantized model look sane.
print("========== SAMPLE GENERATION ==============")
input_ids = processor(text="Hello my name is", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_new_tokens=20)
print(processor.decode(output[0]))
print("==========================================")

qwen.py

import os

import torch
from compressed_tensors.quantization import (
    QuantizationArgs,
    QuantizationScheme,
    QuantizationStrategy,
    QuantizationType,
)
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration

from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.transformers import oneshot

# Load model.
model_id = "Qwen/Qwen2-VL-2B-Instruct"
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_id, device_map="auto", torch_dtype="auto"
)
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

# Oneshot arguments
DATASET_ID = "flickr30k"
DATASET_SPLIT = "test[:3]"
NUM_CALIBRATION_SAMPLES = 1
MAX_SEQUENCE_LENGTH = 2048


# TODO: define real collators in utils
def data_collator(batch):
    assert len(batch) == 1
    return {
        "input_ids": torch.LongTensor(batch[0]["input_ids"]),
        "attention_mask": torch.tensor(batch[0]["attention_mask"]),
        "pixel_values": torch.tensor(
            batch[0]["pixel_values"]
        ),  # torch.Size([14308, 1176])
        "image_grid_thw": torch.tensor(batch[0]["image_grid_thw"]),
    }


# Recipe
recipe = GPTQModifier(
    targets="Linear",
    config_groups={
        "config_group": QuantizationScheme(
            targets=["Linear"],
            weights=QuantizationArgs(
                num_bits=4,
                type=QuantizationType.INT,
                strategy=QuantizationStrategy.GROUP,
                group_size=128,
                symmetric=True,
                dynamic=False,
                actorder="dynamic",
            ),
        ),
    },
    ignore=["re:.*lm_head"],
    update_size=NUM_CALIBRATION_SAMPLES,
    dampening_frac=0.5,
)

# Perform oneshot
save_name = model_id.split("/")[1] + "-W8A8"
save_path = os.path.join("./my_test/", save_name)
print("Starting quantization")
oneshot(
    model=model,
    tokenizer=model_id,
    # dataset=ds,
    dataset=DATASET_ID,
    splits=DATASET_SPLIT,
    recipe=recipe,
    max_seq_length=MAX_SEQUENCE_LENGTH,
    num_calibration_samples=NUM_CALIBRATION_SAMPLES,
    trust_remote_code_model=True,
    output_dir=save_path,
    data_collator=data_collator,
)

processor.save_pretrained(save_path)

# Confirm generations of the quantized model look sane.
print("========== SAMPLE GENERATION ==============")
input_ids = processor(text="Hello my name is", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids, max_new_tokens=20)
print(processor.decode(output[0]))
print("==========================================")

[kylesayrs]

neuralmagic/sparseml#2146

[dsikka]

I think overall looks good. Mostly just a lot of questions.

I would test this against a staged workflow e.g: https://github.com/vllm-project/llm-compressor/tree/main/examples/quantization_2of4_sparse_w4a16
And I think general feedback is minimizing how much we're abstracting away from users.

[dsikka]

LGTM. Why are there ???? in the logs?

[kylesayrs]

@dsikka I use gnu screen which sometimes doesn't handle unicode characters properly. This is a purely a visual bug on my side.

[dsikka]

Approved pending test failures

[kylesayrs]

Tests now pass

[dsikka]

Thanks!

[mgoin]

Thanks for the discussion!