WIP: Multimodal Audio

Dockerfile

@@ -4,6 +4,7 @@ FROM nvcr.io/nvidia/pytorch:24.11-py3
 # Install dependencies.
 RUN apt-get update \
     && apt-get install --no-install-recommends -y acl git-lfs \
+    # && apt-get install --no-install-recommends -y acl git-lfs libtiff5-dev \
     && rm -rf /var/lib/apt/lists/* \
     && git lfs install
 

fast_llm/data/data/gpt/data.py

@@ -32,6 +32,10 @@ class GPTBatch:
     token_ids: torch.Tensor
     loss_masking_spans: list[torch.Tensor] | None = None
     sequence_lengths: list[torch.Tensor] | None = None
+    images: list[torch.Tensor] | None = None
+    image_positions: list[torch.Tensor] | None = None
+    audio: list[torch.Tensor] | None = None
+    audio_positions: list[torch.Tensor] | None = None
 
 
 def gpt_data_collate_fn(batch: list[GPTSample], sampling_parameters: GPTSamplingParameters) -> GPTBatch:
@@ -42,8 +46,44 @@ def gpt_data_collate_fn(batch: list[GPTSample], sampling_parameters: GPTSampling
         stacked_spans = [torch.from_numpy(sample.loss_masking_spans) for sample in batch]
     if not sampling_parameters.cross_document_attention:
         sequence_lengths = [torch.tensor(sample.sequence_lengths) for sample in batch]
+    has_images = False
+    batch_images = []
+    for sample in batch:
+        if sample.images is not None:
+            batch_images.append([torch.from_numpy(image) for image in sample.images])
+            has_images = True
+        else:
+            batch_images.append([])
+    batch_image_positions = []
+    for sample in batch:
+        if sample.image_positions is not None and len(sample.image_positions) > 0:
+            batch_image_positions.append(torch.from_numpy(sample.image_positions))
+        else:
+            batch_image_positions.append([])
+
+    has_audio = False
+    batch_audio = []
+    for sample in batch:
+        if sample.audio is not None and sample.audio_positions is not None:
+            batch_audio.append([torch.from_numpy(audio) for audio in sample.audio])
+            has_audio = True
+        else:
+            batch_audio.append(None)
+    batch_audio_positions = []
+    for sample in batch:
+        if sample.audio_positions is not None:
+            batch_audio_positions.append(torch.from_numpy(sample.audio_positions))
+        else:
+            batch_audio_positions.append([])
+
     return GPTBatch(
-        token_ids=torch.from_numpy(stacked_ids), loss_masking_spans=stacked_spans, sequence_lengths=sequence_lengths
+        token_ids=torch.from_numpy(stacked_ids),
+        loss_masking_spans=stacked_spans,
+        sequence_lengths=sequence_lengths,
+        images=batch_images if has_images else None,
+        image_positions=batch_image_positions if has_images else None,
+        audio=batch_audio if has_audio else None,
+        audio_positions=batch_audio_positions if has_audio else None,
     )
 
 

fast_llm/data/dataset/gpt/config.py

@@ -74,6 +74,15 @@ class GPTSamplingParameters(SamplingParameters):
     vocab_size: int
     use_loss_masking_spans: bool = False
     cross_document_attention: bool = True
+    patch_size: int | None = None
+    image_size: int | None = None
+    aud_downsampling_k: int | None = None
+    aud_padding_duration: int | None = None
+    aud_sampling_rate: int | None = None
+    image_break_token: int | None = None
+    image_end_token: int | None = None
+    audio_start_token: int | None = None
+    audio_end_token: int | None = None
     # How many extra tokens to add to the sequence length.
     # This is used to provide labels even for the last tokens in the sequence.
     extra_tokens: int = 1
@@ -195,11 +204,23 @@ class GPTMemmapDatasetConfig(GPTIndexedDatasetConfig):
         desc="Expected number of tokens in the dataset.",
         hint=FieldHint.optional,
     )
+    num_pixels: int | None = Field(
+        default=None,
+        desc="Expected number of pixels in the dataset.",
+        hint=FieldHint.optional,
+    )
+    num_audio: int | None = Field(
+        default=None,
+        desc="Expected number of audio in the dataset.",
+        hint=FieldHint.optional,
+    )
 
     def build(self) -> "GPTMemmapDataset":
         from fast_llm.data.dataset.gpt.memmap import GPTMemmapDataset
 
-        return GPTMemmapDataset(str(self.path).replace("/", "__"), self.path, self.num_documents, self.num_tokens)
+        return GPTMemmapDataset(
+            str(self.path).replace("/", "__"), self.path, self.num_documents, self.num_tokens, self.num_pixels
+        )
 
 
 @config_class()

fast_llm/data/dataset/gpt/indexed.py

@@ -44,11 +44,20 @@ class GPTDatasetSlice[IndexedDatasetType: GPTIndexedDataset](DatasetSlice[Indexe
 
     def get_document_sizes(self) -> np.ndarray:
         # TODO: This can be really big.
-        return self._dataset.get_document_sizes()[self._begin : self._end]
+        doc_sizes, im_sizes, aud_sizes = self._dataset.get_document_sizes()
+        return (
+            doc_sizes[self._begin : self._end],
+            im_sizes[self._begin : self._end] if im_sizes else [],
+            aud_sizes[self._begin : self._end] if aud_sizes else [],
+        )
 
     def get_document_size(self, index: int) -> int:
         return self._dataset.get_document_size(self._begin + index)
 
+    @property
+    def has_images(self) -> bool:
+        return self._dataset.has_images
+
 
 class GPTConcatenatedDataset[IndexedDatasetType: GPTIndexedDataset](
     ConcatenatedDataset[IndexedDatasetType], GPTIndexedDataset