WIP: truncate in batch

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

examples/multimodal_vision/gemma3_example.py

@@ -25,9 +25,13 @@
 
 # Define a oneshot data collator for multimodal processors
 # remove extra dim added by vision processor
-def data_collator(features: list[dict[str, object]]):
-    features = [{key: feature[key][0] for key in feature} for feature in features]
-    return collator(features)
+# def data_collator(features: list[dict[str, object]]):
+#     features = [{key: feature[key][0] for key in feature} for feature in features]
+#     return collator(features)
+# Define a oneshot data collator for multimodal inputs.
+def data_collator(batch):
+    assert len(batch) == 1
+    return {key: torch.tensor(value) for key, value in batch[0].items()}
 
 
 # Recipe
@@ -57,10 +61,10 @@ def data_collator(features: list[dict[str, object]]):
         max_seq_length=MAX_SEQUENCE_LENGTH,
         num_calibration_samples=NUM_CALIBRATION_SAMPLES,
         data_collator=data_collator,
-        trust_remote_code_model=True,
         pipeline="sequential",
     )
 import torch
+
 del prof._memory.timeline[torch.device("cpu")]
 prof.save_memory_timeline("with_disable.png")
 exit(0)

examples/quantization_w4a16/llama3_example.py

@@ -7,7 +7,7 @@
 
 # Select model and load it.
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-#model_id = "meta-llama/Llama-3.2-1B-Instruct"
+# model_id = "meta-llama/Llama-3.2-1B-Instruct"
 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
@@ -64,7 +64,7 @@ def tokenize(sample):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     pipeline="sequential",
-    shuffle_calibration_samples=False
+    shuffle_calibration_samples=False,
 )
 exit(0)
 

src/llmcompressor/args/dataset_arguments.py

@@ -8,7 +8,9 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Callable, Optional
+from typing import Callable
+
+from loguru import logger
 
 
 @dataclass
@@ -67,15 +69,16 @@ class CustomDatasetArguments(DVCDatasetArguments):
         },
     )
 
-    data_collator: Optional[Callable] = field(
-        default=None,
+    data_collator: str | Callable = field(
+        default="truncation",
         metadata={
             "help": (
                 "The function to used to form a batch from the dataset. Defaults to "
                 "`DataCollatorWithPadding(processor)`."
             )
         },
     )
+    # remove_extra_processor_dim: bool = field
 
     batch_size: int = field(
         default=1,
@@ -154,7 +157,7 @@ class DatasetArguments(CustomDatasetArguments):
         default=False,
         metadata={"help": "Overwrite the cached preprocessed datasets or not."},
     )
-    num_data_workers: int | None = field(
+    preprocessing_num_workers: int | None = field(
         default=None,
         metadata={"help": "The number of workers to use for dataset processing."},
     )
@@ -240,3 +243,17 @@ class DatasetArguments(CustomDatasetArguments):
 
     def is_dataset_provided(self) -> bool:
         return self.dataset is not None or self.dataset_path is not None
+
+    def get_num_data_workers(self):
+        import multiprocessing
+
+        if self.preprocessing_num_workers is not None:
+            return self.preprocessing_num_workers
+
+        try:
+            return min(multiprocessing.cpu_count() // 2, 8)  # cap max at 8
+        except NotImplementedError:
+            logger.warning(
+                "Could not determine number of CPUs, defaulting to 1 dataloader worker."
+            )
+            return 1

src/llmcompressor/datasets/utils.py

@@ -7,18 +7,15 @@
 one-shot calibration workflows.
 """
 
-import math
-import multiprocessing
 import re
-from typing import Any, Callable, Optional
 from collections.abc import Iterator, Sized
-from torch.utils._pytree import tree_flatten
+from typing import Any, Callable, Optional
 
 import torch
 from datasets import Dataset
 from loguru import logger
-from torch.utils.data import DataLoader, SequentialSampler, RandomSampler, Sampler
-from transformers.data import DataCollatorWithPadding
+from torch.utils.data import DataLoader, RandomSampler, Sampler
+from transformers.data import DataCollatorWithPadding, default_data_collator
 
 from llmcompressor.args import DatasetArguments
 from llmcompressor.transformers.data import TextGenerationDataset
@@ -110,71 +107,31 @@ def get_calibration_dataloader(
         # weight-only quantization or dynamic quantization
         return
 
+    # load and tokenize dataset
     datasets = get_processed_dataset(
         dataset_args=dataset_args,
         processor=processor,
         do_oneshot=True,
         do_train=False,
     )
-
     calibration_dataset = datasets.get("calibration")
-    collate_fn = dataset_args.data_collator or _make_padding_collator(processor)
-    num_workers = dataset_args.num_data_workers or _infer_num_data_workers()
-
-    return format_calibration_data(
-        tokenized_dataset=calibration_dataset,
-        collate_fn=collate_fn,
-        batch_size=dataset_args.batch_size,
-        num_calibration_samples=dataset_args.num_calibration_samples,
-        do_shuffle=dataset_args.shuffle_calibration_samples,
-        num_workers=num_workers,
-    )
+
+    return format_calibration_data(dataset_args, calibration_dataset, processor)
 
 
 def format_calibration_data(
+    args: DatasetArguments,
     tokenized_dataset: Dataset,
-    collate_fn: Callable,
-    batch_size: int = 1,
-    num_calibration_samples: int | None = None,
-    do_shuffle: bool = False,
-    num_workers: int = 1,
+    processor: Processor,
 ) -> list[torch.Tensor]:
-    """
-    Creates a dataloader out of the calibration dataset split, trimming it to
-    the desired number of calibration samples
-    :param tokenized_dataset: dataset to convert to dataloader
-    :param num_calibration_samples: number of batches to convert
-    :param do_shuffle: whether to shuffle the dataset before selecting calibration
-        samples, true by default
-    :param collate_fn: optional custom collate function, or use default
-    :return: list of trimmed calibration data tensors
-    """
-    # (1) shuffle before truncating
-    if do_shuffle:
-        tokenized_dataset = tokenized_dataset.shuffle()
-
-    # (2) truncate dataset
-    if num_calibration_samples is not None:
-        if num_calibration_samples > len(tokenized_dataset):
-            logger.warning(
-                f"Requested {num_calibration_samples} calibration samples but the "
-                f"provided dataset only has {len(tokenized_dataset)} samples."
-            )
-            num_calibration_samples = len(tokenized_dataset)
-        tokenized_dataset = tokenized_dataset.select(range(num_calibration_samples))
-
-    # (3) create sampler
-    sampler = _make_sampler(tokenized_dataset, num_calibration_samples, do_shuffle)
-
-    # (4) create dataloader
-    dataloader_params = {
-        "batch_size": batch_size,
-        "sampler": sampler,
-        "collate_fn": collate_fn,
-        "pin_memory": False,
-        "num_workers": num_workers,
-    }
-    return DataLoader(tokenized_dataset, **dataloader_params)
+    return DataLoader(
+        tokenized_dataset,
+        batch_size=args.batch_size,
+        sampler=_make_sampler(args, tokenized_dataset),
+        collate_fn=_make_collate_fn(args, processor),
+        pin_memory=False,
+        num_workers=args.get_num_data_workers(),
+    )
 
 
 def make_dataset_splits(
@@ -216,25 +173,30 @@ def make_dataset_splits(
     return split_datasets
 
 
-def _make_padding_collator(processor: Processor) -> DataCollatorWithPadding:
-    tokenizer = getattr(processor, "tokenizer", processor)
-    if tokenizer.pad_token is None or tokenizer.pad_token_id < 0:
-        logger.debug("Could not find padding token. Setting PAD token to EOS token")
-        tokenizer.pad_token = tokenizer.eos_token
+def _make_collate_fn(dataset_args: DatasetArguments, processor: Processor) -> Callable:
+    if isinstance(dataset_args.data_collator, Callable):
+        return dataset_args.data_collator
 
-    return DataCollatorWithPadding(tokenizer)
+    if dataset_args.data_collator == "truncation":
+        return data_collator_with_truncation
+
+    elif dataset_args.data_collator == "padding":
+        tokenizer = getattr(processor, "tokenizer", processor)
+        if tokenizer.pad_token is None or tokenizer.pad_token_id < 0:
+            logger.debug("Could not find padding token. Setting PAD token to EOS token")
+            tokenizer.pad_token = tokenizer.eos_token
+
+        return DataCollatorWithPadding(tokenizer)
+
+    else:
+        assert False
+
+
+def _make_sampler(args: DatasetArguments, dataset: Dataset) -> Sampler:
+    num_samples = args.num_calibration_samples
+    shuffle = args.shuffle_calibration_samples
+    batch_size = args.batch_size
 
-def _infer_num_data_workers() -> int:
-    MAX_DATALOADER_WORKERS = 8
-    try:
-        return min(MAX_DATALOADER_WORKERS, multiprocessing.cpu_count() // 2)
-    except NotImplementedError:
-        logger.warning(
-            "Could not determine number of CPUs, defaulting to 0 dataloader workers."
-        )
-        return 0
-    
-def _make_sampler(dataset: Dataset, num_samples: int | None, shuffle: bool) -> Sampler:
     if num_samples is not None and num_samples > len(dataset):
         logger.warning(
             f"Requested {num_samples} samples but the provided dataset only has "
@@ -243,19 +205,38 @@ def _make_sampler(dataset: Dataset, num_samples: int | None, shuffle: bool) -> S
         num_samples = len(dataset)
 
     if shuffle:
-        return RandomSampler(
-            dataset,
-            replacement=False,
-            num_samples=num_samples,
-            generator=None,
-        )
+        if batch_size > 1:
+            logger.warning(
+                "Shuffling a dataset can lead to unoptimal batching for sequence "
+                "lengths non-uniform sizes. When collating with truncation, this will "
+                "delete a large number of tokens. When collating with padding, this "
+                "will add a large number of padding tokens.\n\nPlease consider calling "
+                "`oneshot` with `batch_size=1`"
+            )
+
+        return RandomSampler(dataset, num_samples=num_samples)
     else:
-        return LengthAwareSampler(
-            dataset,
-            replacement=False,
-            num_samples=num_samples,
-            generator=None,
-        )
+        return LengthAwareSampler(dataset, num_samples=num_samples)
+
+
+def data_collator_with_truncation(
+    features: list[dict[str, Any]], return_tensors: str = "pt"
+) -> dict[str, Any]:
+    total_removed = 0
+    total = 0
+
+    keys = set().union(*(feature.keys() for feature in features))
+    for key in keys:
+        lengths = [
+            len(feature[key]) for feature in features if isinstance(feature[key], list)
+        ]
+        min_len = min(lengths)
+        for feature in features:
+            total_removed += len(feature[key]) - min_len
+            total += len(feature[key])
+            feature[key] = feature[key][:min_len]
+
+    return default_data_collator(features, return_tensors)
 
 
 class LengthAwareSampler(Sampler[int]):
@@ -265,20 +246,11 @@ class LengthAwareSampler(Sampler[int]):
     def __init__(
         self,
         data_source: Sized,
-        replacement: bool = False,
         num_samples: Optional[int] = None,
-        generator: Optional[torch.Generator] = None,
     ) -> None:
         self.data_source = data_source
-        self.replacement = replacement
         self._num_samples = num_samples
 
-        if replacement:
-            raise NotImplementedError()
-
-        if generator:
-            raise NotImplementedError()
-
         lengths = [len(sample) for sample in data_source["input_ids"]]
         self.order = torch.argsort(torch.tensor(lengths), descending=True).tolist()
 
@@ -293,4 +265,4 @@ def __iter__(self) -> Iterator[int]:
         return iter(self.order)
 
     def __len__(self) -> int:
-        return self._num_samples
+        return self._num_samples

src/llmcompressor/entrypoints/oneshot.py

@@ -249,7 +249,7 @@ def oneshot(
     dataset_path: str | None = None,
     splits: str | list[str] | dict[str, str] | None = None,
     batch_size: int = 1,
-    data_collator: Optional[Callable] = None,
+    data_collator: Optional[Callable] = "truncation",
     num_calibration_samples: int = 512,
     shuffle_calibration_samples: bool = True,
     max_seq_length: int = 384,
@@ -258,7 +258,7 @@ def oneshot(
     concatenate_data: bool = False,
     streaming: bool = False,
     overwrite_cache: bool = False,
-    num_data_workers: int | None = None,
+    preprocessing_num_workers: int | None = None,
     min_tokens_per_module: float | None = None,
     moe_calibrate_all_experts: bool = True,
     quantization_aware_calibration: bool = True,
@@ -319,7 +319,7 @@ def oneshot(
         max_seq_length.
     :param streaming: True to stream data from a cloud dataset.
     :param overwrite_cache: Whether to overwrite the cached preprocessed datasets.
-    :param num_data_workers: Number of processes for dataset preprocessing.
+    :param preprocessing_num_workers: Number of processes for dataset preprocessing.
     :param min_tokens_per_module: Minimum percentage of tokens per
         module, relevant for MoE models.
     :param moe_calibrate_all_experts: Whether to calibrate all experts during MoE